"""Tests for _baseline — JIT detection and environment configuration."""

from __future__ import annotations

import textwrap
from typing import TYPE_CHECKING
from unittest.mock import MagicMock, patch

import attrs
import pytest

if TYPE_CHECKING:
    from pathlib import Path

from codeflash_python.test_discovery.models import TestType
from codeflash_python.testing.models import (
    FunctionTestInvocation,
    InvocationId,
    TestConfig,
    TestFile,
    TestFiles,
    TestResults,
)
from codeflash_python.verification._baseline import (
    JIT_DECORATORS,
    contains_jit_decorator,
    establish_original_code_baseline,
    jit_disabled_env,
)
from codeflash_python.verification.models import OriginalCodeBaseline


class TestContainsJitDecorator:
    """contains_jit_decorator JIT decorator detection."""

    def test_numba_jit(self) -> None:
        """Detects @numba.jit decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_njit(self) -> None:
        """Detects @numba.njit decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.njit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_vectorize(self) -> None:
        """Detects @numba.vectorize decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.vectorize
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_guvectorize(self) -> None:
        """Detects @numba.guvectorize decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.guvectorize
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_stencil(self) -> None:
        """Detects @numba.stencil decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.stencil
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_cfunc(self) -> None:
        """Detects @numba.cfunc decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.cfunc
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_generated_jit(self) -> None:
        """Detects @numba.generated_jit decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.generated_jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_alias(self) -> None:
        """Detects JIT decorator when numba is aliased."""
        code = textwrap.dedent("""\
            import numba as nb

            @nb.jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_from_import(self) -> None:
        """Detects JIT decorator via from-import."""
        code = textwrap.dedent("""\
            from numba import jit

            @jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_from_import_alias(self) -> None:
        """Detects JIT decorator via aliased from-import."""
        code = textwrap.dedent("""\
            from numba import jit as my_jit

            @my_jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_with_arguments(self) -> None:
        """Detects @jit(nopython=True) decorator with arguments."""
        code = textwrap.dedent("""\
            from numba import jit

            @jit(nopython=True)
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_cuda_jit(self) -> None:
        """Detects @numba.cuda.jit decorator."""
        code = textwrap.dedent("""\
            import numba

            @numba.cuda.jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_numba_cuda_from_import(self) -> None:
        """Detects @jit via from numba.cuda import jit."""
        code = textwrap.dedent("""\
            from numba.cuda import jit

            @jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_torch_compile(self) -> None:
        """Detects @torch.compile decorator."""
        code = textwrap.dedent("""\
            import torch

            @torch.compile
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_torch_jit_script(self) -> None:
        """Detects @torch.jit.script decorator."""
        code = textwrap.dedent("""\
            import torch

            @torch.jit.script
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_torch_jit_trace(self) -> None:
        """Detects @torch.jit.trace decorator."""
        code = textwrap.dedent("""\
            import torch

            @torch.jit.trace
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_jax_jit(self) -> None:
        """Detects @jax.jit decorator."""
        code = textwrap.dedent("""\
            import jax

            @jax.jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_jax_from_import(self) -> None:
        """Detects JIT decorator via from jax import jit."""
        code = textwrap.dedent("""\
            from jax import jit

            @jit
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_tensorflow_function(self) -> None:
        """Detects @tensorflow.function decorator."""
        code = textwrap.dedent("""\
            import tensorflow

            @tensorflow.function
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_tensorflow_alias(self) -> None:
        """Detects @tf.function when tensorflow is aliased as tf."""
        code = textwrap.dedent("""\
            import tensorflow as tf

            @tf.function
            def f():
                pass
        """)
        assert contains_jit_decorator(code) is True

    def test_non_jit_decorators(self) -> None:
        """Returns False for standard non-JIT decorators."""
        code = textwrap.dedent("""\
            class Foo:
                @property
                def bar(self):
                    return 1

                @staticmethod
                def baz():
                    pass

            @my_decorator
            def qux():
                pass
        """)
        assert contains_jit_decorator(code) is False

    def test_no_decorators(self) -> None:
        """Returns False when no decorators are present."""
        code = textwrap.dedent("""\
            def f():
                pass

            def g():
                return 1
        """)
        assert contains_jit_decorator(code) is False

    def test_syntax_error(self) -> None:
        """Returns False for code with syntax errors."""
        code = "def f(\n    not valid !!!"
        assert contains_jit_decorator(code) is False

    def test_empty_string(self) -> None:
        """Returns False for empty string input."""
        assert contains_jit_decorator("") is False

    def test_imports_but_no_functions(self) -> None:
        """Returns False when JIT modules imported but no functions."""
        code = textwrap.dedent("""\
            import numba
            import torch

            x = 1
        """)
        assert contains_jit_decorator(code) is False


class TestJitDisabledEnv:
    """jit_disabled_env environment variable dictionary."""

    def test_returns_dict_with_correct_keys(self) -> None:
        """Returns dict containing all expected JIT-disabling env var keys."""
        result = jit_disabled_env()
        expected_keys = {
            "NUMBA_DISABLE_JIT",
            "TORCHDYNAMO_DISABLE",
            "PYTORCH_JIT",
            "TF_XLA_FLAGS",
            "TF_ENABLE_ONEDNN_OPTS",
            "JAX_DISABLE_JIT",
        }
        assert expected_keys == set(result.keys())

    def test_returns_dict_with_string_values(self) -> None:
        """All values in the returned dict are strings."""
        result = jit_disabled_env()
        assert all(isinstance(v, str) for v in result.values())

    def test_returns_dict(self) -> None:
        """Return type is a dict."""
        result = jit_disabled_env()
        assert isinstance(result, dict)


class TestJitDecorators:
    """JIT_DECORATORS constant structure."""

    def test_contains_expected_modules(self) -> None:
        """Contains entries for all supported JIT modules."""
        expected_modules = {
            "numba",
            "numba.cuda",
            "torch",
            "torch.jit",
            "tensorflow",
            "jax",
        }
        assert expected_modules == set(JIT_DECORATORS.keys())

    def test_numba_has_expected_decorators(self) -> None:
        """Numba entry contains all expected decorator names."""
        expected = {
            "jit",
            "njit",
            "vectorize",
            "guvectorize",
            "stencil",
            "cfunc",
            "generated_jit",
        }
        assert expected == JIT_DECORATORS["numba"]

    def test_values_are_sets_of_strings(self) -> None:
        """All values are sets containing only strings."""
        for module, decorators in JIT_DECORATORS.items():
            assert isinstance(decorators, set), f"{module} value is not a set"
            assert all(isinstance(d, str) for d in decorators), (
                f"{module} contains non-string decorator"
            )


def _make_test_results(
    file_name: Path,
    *,
    runtime: int = 1000,
    did_pass: bool = True,
    test_type: TestType = TestType.EXISTING_UNIT_TEST,
) -> TestResults:
    """Build a TestResults with one invocation."""
    results = TestResults()
    results.add(
        FunctionTestInvocation(
            loop_index=1,
            id=InvocationId(
                test_module_path="test_module",
                test_class_name=None,
                test_function_name="test_func",
                function_getting_tested="target_func",
                iteration_id="0",
            ),
            file_name=file_name,
            did_pass=did_pass,
            runtime=runtime,
            test_framework="pytest",
            test_type=test_type,
            return_value=None,
            timed_out=False,
        ),
    )
    return results


class TestOriginalCodeBaseline:
    """OriginalCodeBaseline frozen data class."""

    def test_construction(self, tmp_path: Path) -> None:
        """Can be constructed with all required fields."""
        tr = _make_test_results(tmp_path / "test.py")
        baseline = OriginalCodeBaseline(
            behavior_test_results=tr,
            benchmarking_test_results=tr,
            runtime=5000,
            line_profile_results=tr,
        )
        assert tr is baseline.behavior_test_results
        assert 5000 == baseline.runtime

    def test_default_functions_to_remove(self, tmp_path: Path) -> None:
        """Default functions_to_remove is an empty tuple."""
        tr = _make_test_results(tmp_path / "test.py")
        baseline = OriginalCodeBaseline(
            behavior_test_results=tr,
            benchmarking_test_results=tr,
            runtime=5000,
            line_profile_results=tr,
        )
        assert () == baseline.functions_to_remove

    def test_custom_functions_to_remove(self, tmp_path: Path) -> None:
        """Accepts explicit functions_to_remove tuple."""
        tr = _make_test_results(tmp_path / "test.py")
        baseline = OriginalCodeBaseline(
            behavior_test_results=tr,
            benchmarking_test_results=tr,
            runtime=5000,
            line_profile_results=tr,
            functions_to_remove=("fn_a", "fn_b"),
        )
        assert ("fn_a", "fn_b") == baseline.functions_to_remove

    def test_frozen(self, tmp_path: Path) -> None:
        """Raises on attribute assignment (frozen)."""
        tr = _make_test_results(tmp_path / "test.py")
        baseline = OriginalCodeBaseline(
            behavior_test_results=tr,
            benchmarking_test_results=tr,
            runtime=5000,
            line_profile_results=tr,
        )
        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            baseline.runtime = 9999  # type: ignore[misc]

    def test_field_access(self, tmp_path: Path) -> None:
        """All fields are accessible after construction."""
        behavioral = _make_test_results(tmp_path / "test.py")
        benchmarking = _make_test_results(tmp_path / "test.py", runtime=2000)
        line_profile = _make_test_results(tmp_path / "test.py", runtime=3000)
        baseline = OriginalCodeBaseline(
            behavior_test_results=behavioral,
            benchmarking_test_results=benchmarking,
            runtime=7000,
            line_profile_results=line_profile,
            functions_to_remove=("fn_x",),
        )
        assert behavioral is baseline.behavior_test_results
        assert benchmarking is baseline.benchmarking_test_results
        assert 7000 == baseline.runtime
        assert line_profile is baseline.line_profile_results
        assert ("fn_x",) == baseline.functions_to_remove


class TestEstablishOriginalCodeBaseline:
    """establish_original_code_baseline orchestration."""

    def _make_fixtures(
        self,
        tmp_path: Path,
    ) -> tuple[TestFiles, TestConfig, dict[str, str]]:
        """Build common test fixtures."""
        test_files = TestFiles(
            test_files=[
                TestFile(
                    original_file_path=tmp_path / "test_example.py",
                    instrumented_behavior_file_path=(
                        tmp_path / "test_example_behavior.py"
                    ),
                    benchmarking_file_path=(
                        tmp_path / "test_example_bench.py"
                    ),
                ),
            ],
        )
        test_config = TestConfig(
            tests_project_rootdir=tmp_path,
        )
        test_env: dict[str, str] = {"PATH": "/usr/bin"}
        return test_files, test_config, test_env

    @patch("codeflash_python.testing._parse_results.parse_test_results")
    @patch(
        "codeflash_python.testing._test_runner.async_run_line_profile_tests"
    )
    @patch(
        "codeflash_python.testing._test_runner.async_run_benchmarking_tests"
    )
    @patch("codeflash_python.testing._test_runner.async_run_behavioral_tests")
    async def test_successful_baseline(
        self,
        mock_run_behavioral: MagicMock,
        mock_run_benchmarking: MagicMock,
        mock_run_line_profile: MagicMock,
        mock_parse_results: MagicMock,
        tmp_path: Path,
    ) -> None:
        """Returns OriginalCodeBaseline with correct fields on success."""
        test_files, test_config, test_env = self._make_fixtures(
            tmp_path,
        )
        test_file = tmp_path / "test.py"
        xml_path = tmp_path / "results.xml"

        behavioral_results = _make_test_results(test_file, runtime=1000)
        line_profile_results = _make_test_results(test_file, runtime=500)
        benchmarking_results = _make_test_results(test_file, runtime=2000)

        mock_run_behavioral.return_value = (
            xml_path,
            MagicMock(),
            None,
            None,
        )
        mock_run_benchmarking.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_run_line_profile.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_parse_results.side_effect = [
            behavioral_results,
            line_profile_results,
            benchmarking_results,
        ]

        result = await establish_original_code_baseline(
            test_files=test_files,
            test_config=test_config,
            test_env=test_env,
            cwd=tmp_path,
        )

        assert result is not None
        assert isinstance(result, OriginalCodeBaseline)
        assert behavioral_results is result.behavior_test_results
        assert benchmarking_results is result.benchmarking_test_results
        assert line_profile_results is result.line_profile_results
        assert result.runtime > 0

    @patch("codeflash_python.testing._parse_results.parse_test_results")
    @patch(
        "codeflash_python.testing._test_runner.async_run_line_profile_tests"
    )
    @patch(
        "codeflash_python.testing._test_runner.async_run_benchmarking_tests"
    )
    @patch("codeflash_python.testing._test_runner.async_run_behavioral_tests")
    async def test_empty_behavioral_returns_none(
        self,
        mock_run_behavioral: MagicMock,
        mock_run_benchmarking: MagicMock,
        mock_run_line_profile: MagicMock,
        mock_parse_results: MagicMock,
        tmp_path: Path,
    ) -> None:
        """Returns None when behavioral results are empty."""
        test_files, test_config, test_env = self._make_fixtures(
            tmp_path,
        )

        mock_run_behavioral.return_value = (
            tmp_path / "results.xml",
            MagicMock(),
            None,
            None,
        )
        mock_parse_results.return_value = TestResults()

        result = await establish_original_code_baseline(
            test_files=test_files,
            test_config=test_config,
            test_env=test_env,
            cwd=tmp_path,
        )

        assert result is None

    @patch("codeflash_python.testing._parse_results.parse_test_results")
    @patch(
        "codeflash_python.testing._test_runner.async_run_line_profile_tests"
    )
    @patch(
        "codeflash_python.testing._test_runner.async_run_benchmarking_tests"
    )
    @patch("codeflash_python.testing._test_runner.async_run_behavioral_tests")
    async def test_zero_benchmark_runtime_returns_none(
        self,
        mock_run_behavioral: MagicMock,
        mock_run_benchmarking: MagicMock,
        mock_run_line_profile: MagicMock,
        mock_parse_results: MagicMock,
        tmp_path: Path,
    ) -> None:
        """Returns None when benchmark runtime is zero."""
        test_files, test_config, test_env = self._make_fixtures(
            tmp_path,
        )
        test_file = tmp_path / "test.py"
        xml_path = tmp_path / "results.xml"

        behavioral_results = _make_test_results(test_file, runtime=1000)
        line_profile_results = _make_test_results(test_file, runtime=500)
        zero_benchmarking = _make_test_results(
            test_file,
            runtime=0,
            did_pass=False,
        )

        mock_run_behavioral.return_value = (
            xml_path,
            MagicMock(),
            None,
            None,
        )
        mock_run_benchmarking.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_run_line_profile.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_parse_results.side_effect = [
            behavioral_results,
            line_profile_results,
            zero_benchmarking,
        ]

        result = await establish_original_code_baseline(
            test_files=test_files,
            test_config=test_config,
            test_env=test_env,
            cwd=tmp_path,
        )

        assert result is None

    @patch("codeflash_python.testing._parse_results.parse_test_results")
    @patch(
        "codeflash_python.testing._test_runner.async_run_line_profile_tests"
    )
    @patch(
        "codeflash_python.testing._test_runner.async_run_benchmarking_tests"
    )
    @patch("codeflash_python.testing._test_runner.async_run_behavioral_tests")
    async def test_precomputed_behavioral_skips_behavioral_run(
        self,
        mock_run_behavioral: MagicMock,
        mock_run_benchmarking: MagicMock,
        mock_run_line_profile: MagicMock,
        mock_parse_results: MagicMock,
        tmp_path: Path,
    ) -> None:
        """Skips running behavioral tests when precomputed_behavioral given."""
        test_files, test_config, test_env = self._make_fixtures(
            tmp_path,
        )
        test_file = tmp_path / "test.py"
        xml_path = tmp_path / "results.xml"

        precomputed = _make_test_results(test_file, runtime=1000)
        line_profile_results = _make_test_results(test_file, runtime=500)
        benchmarking_results = _make_test_results(test_file, runtime=2000)

        mock_run_benchmarking.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_run_line_profile.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_parse_results.side_effect = [
            line_profile_results,
            benchmarking_results,
        ]

        result = await establish_original_code_baseline(
            test_files=test_files,
            test_config=test_config,
            test_env=test_env,
            cwd=tmp_path,
            precomputed_behavioral=precomputed,
        )

        mock_run_behavioral.assert_not_called()
        assert result is not None
        assert precomputed is result.behavior_test_results

    @patch("codeflash_python.testing._parse_results.parse_test_results")
    @patch(
        "codeflash_python.testing._test_runner.async_run_line_profile_tests"
    )
    @patch(
        "codeflash_python.testing._test_runner.async_run_benchmarking_tests"
    )
    @patch("codeflash_python.testing._test_runner.async_run_behavioral_tests")
    async def test_failed_regression_in_functions_to_remove(
        self,
        mock_run_behavioral: MagicMock,
        mock_run_benchmarking: MagicMock,
        mock_run_line_profile: MagicMock,
        mock_parse_results: MagicMock,
        tmp_path: Path,
    ) -> None:
        """Failed GENERATED_REGRESSION tests appear in functions_to_remove."""
        test_files, test_config, test_env = self._make_fixtures(
            tmp_path,
        )
        test_file = tmp_path / "test.py"
        xml_path = tmp_path / "results.xml"

        behavioral_results = TestResults()
        behavioral_results.add(
            FunctionTestInvocation(
                loop_index=1,
                id=InvocationId(
                    test_module_path="test_module",
                    test_class_name=None,
                    test_function_name="test_passing",
                    function_getting_tested="target_func",
                    iteration_id="0",
                ),
                file_name=test_file,
                did_pass=True,
                runtime=1000,
                test_framework="pytest",
                test_type=TestType.EXISTING_UNIT_TEST,
                return_value=None,
                timed_out=False,
            ),
        )
        behavioral_results.add(
            FunctionTestInvocation(
                loop_index=1,
                id=InvocationId(
                    test_module_path="test_module",
                    test_class_name=None,
                    test_function_name="test_generated_fail",
                    function_getting_tested="target_func",
                    iteration_id="1",
                ),
                file_name=test_file,
                did_pass=False,
                runtime=500,
                test_framework="pytest",
                test_type=TestType.GENERATED_REGRESSION,
                return_value=None,
                timed_out=False,
            ),
        )

        line_profile_results = _make_test_results(test_file, runtime=500)
        benchmarking_results = _make_test_results(test_file, runtime=2000)

        mock_run_behavioral.return_value = (
            xml_path,
            MagicMock(),
            None,
            None,
        )
        mock_run_benchmarking.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_run_line_profile.return_value = (
            xml_path,
            MagicMock(),
        )
        mock_parse_results.side_effect = [
            behavioral_results,
            line_profile_results,
            benchmarking_results,
        ]

        result = await establish_original_code_baseline(
            test_files=test_files,
            test_config=test_config,
            test_env=test_env,
            cwd=tmp_path,
        )

        assert result is not None
        assert "test_generated_fail" in result.functions_to_remove
