"""Baseline establishment utilities.

Provides JIT compilation detection, environment variable helpers,
and orchestration for establishing baseline metrics on original code.
"""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize
    from ..testing.models import TestConfig, TestFiles, TestResults
    from .models import OriginalCodeBaseline

log = logging.getLogger(__name__)

JIT_DECORATORS: dict[str, set[str]] = {
    "numba": {
        "jit",
        "njit",
        "vectorize",
        "guvectorize",
        "stencil",
        "cfunc",
        "generated_jit",
    },
    "numba.cuda": {"jit"},
    "torch": {"compile"},
    "torch.jit": {"script", "trace"},
    "tensorflow": {"function"},
    "jax": {"jit"},
}


class JitDecoratorDetector(ast.NodeVisitor):
    """AST visitor that detects JIT compilation decorators.

    Tracks import aliases to correctly resolve decorators from
    numba, torch, tensorflow, and jax.
    """

    def __init__(self) -> None:
        """Initialize import alias tracking state."""
        # Maps local name -> (module, original_name)
        # e.g., {"nb": ("numba", None), "my_jit": ("numba", "jit")}
        self.import_aliases: dict[str, tuple[str, str | None]] = {}
        self.found_jit_decorator = False

    def visit_Import(self, node: ast.Import) -> None:
        """Track regular imports like 'import numba'."""
        for alias in node.names:
            local_name = alias.asname or alias.name
            self.import_aliases[local_name] = (alias.name, None)
        self.generic_visit(node)

    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
        """Track from-imports like 'from numba import jit'."""
        if node.module is None:
            self.generic_visit(node)
            return
        for alias in node.names:
            local_name = alias.asname or alias.name
            self.import_aliases[local_name] = (node.module, alias.name)
        self.generic_visit(node)

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        """Check function decorators for JIT decorators."""
        for decorator in node.decorator_list:
            if self._is_jit_decorator(decorator):
                self.found_jit_decorator = True
                return
        self.generic_visit(node)

    def _is_jit_decorator(self, node: ast.expr) -> bool:
        """Check if a decorator node is a known JIT decorator."""
        if isinstance(node, ast.Call):
            return self._is_jit_decorator(node.func)
        if isinstance(node, ast.Name):
            return self._check_name_decorator(node.id)
        if isinstance(node, ast.Attribute):
            return self._check_attribute_decorator(node)
        return False

    def _check_name_decorator(self, name: str) -> bool:
        """Check if a simple name decorator is a JIT decorator."""
        if name not in self.import_aliases:
            return False
        module, imported_name = self.import_aliases[name]
        if imported_name is None:
            return False
        return self._is_known_jit_decorator(module, imported_name)

    def _check_attribute_decorator(self, node: ast.Attribute) -> bool:
        """Check if an attribute decorator is a JIT decorator."""
        parts = self._get_attribute_parts(node)
        if not parts:
            return False
        first_part = parts[0]
        rest_parts = parts[1:]
        if first_part in self.import_aliases:
            module, imported_name = self.import_aliases[first_part]
            if imported_name is None:
                if rest_parts:
                    full_module = module
                    decorator_name = rest_parts[-1]
                    if len(rest_parts) > 1:
                        full_module = f"{module}.{'.'.join(rest_parts[:-1])}"
                    return self._is_known_jit_decorator(
                        full_module, decorator_name
                    )
            elif rest_parts:
                full_module = f"{module}.{imported_name}"
                decorator_name = rest_parts[-1]
                if len(rest_parts) > 1:
                    full_module = f"{full_module}.{'.'.join(rest_parts[:-1])}"
                return self._is_known_jit_decorator(
                    full_module, decorator_name
                )
        elif rest_parts:
            full_module = first_part
            if len(rest_parts) > 1:
                full_module = f"{first_part}.{'.'.join(rest_parts[:-1])}"
            decorator_name = rest_parts[-1]
            return self._is_known_jit_decorator(full_module, decorator_name)
        return False

    def _get_attribute_parts(self, node: ast.Attribute) -> list[str]:
        """Get all parts of an attribute chain."""
        parts: list[str] = []
        current: ast.expr = node
        while isinstance(current, ast.Attribute):
            parts.append(current.attr)
            current = current.value
        if isinstance(current, ast.Name):
            parts.append(current.id)
            parts.reverse()
            return parts
        return []

    def _is_known_jit_decorator(
        self, module: str, decorator_name: str
    ) -> bool:
        """Check if a decorator from a module is a known JIT one."""
        if module in JIT_DECORATORS:
            return decorator_name in JIT_DECORATORS[module]
        return False


def contains_jit_decorator(code: str) -> bool:
    """Check if code contains JIT compilation decorators.

    Detects decorators from numba, torch, tensorflow, and jax,
    handling import aliases and decorator arguments.
    """
    try:
        tree = ast.parse(code)
    except SyntaxError:
        return False
    detector = JitDecoratorDetector()
    detector.visit(tree)
    return detector.found_jit_decorator


def jit_disabled_env() -> dict[str, str]:
    """Return environment variables that disable JIT compilation.

    Used during coverage measurement to prevent JIT compilers from
    interfering with coverage instrumentation.
    """
    return {
        "NUMBA_DISABLE_JIT": "1",
        "TORCHDYNAMO_DISABLE": "1",
        "PYTORCH_JIT": "0",
        "TF_XLA_FLAGS": "--tf_xla_auto_jit=0",
        "TF_ENABLE_ONEDNN_OPTS": "0",
        "JAX_DISABLE_JIT": "0",
    }


async def establish_original_code_baseline(  # noqa: PLR0913
    test_files: TestFiles,
    test_config: TestConfig,
    test_env: dict[str, str],
    cwd: Path,
    optimization_iteration: int = 0,
    precomputed_behavioral: TestResults | None = None,
    is_async: bool = False,  # noqa: FBT001, FBT002
    async_function: FunctionToOptimize | None = None,
) -> OriginalCodeBaseline | None:
    """Orchestrate baseline establishment for the original code.

    Runs behavioral tests, line profiling, and performance benchmarks
    on the original (unoptimized) code.  Returns the complete baseline
    metrics or *None* if tests fail to produce usable results.

    For async functions, *async_function* should be provided so that the
    performance decorator can be applied during benchmarking to capture
    throughput markers.
    """
    from ..test_discovery.models import TestType  # noqa: PLC0415
    from ..testing._parse_results import parse_test_results  # noqa: PLC0415
    from ..testing._test_runner import (  # noqa: PLC0415
        async_run_behavioral_tests,
        async_run_benchmarking_tests,
        async_run_line_profile_tests,
    )
    from .models import OriginalCodeBaseline  # noqa: PLC0415

    # Step 1: behavioral tests (with coverage collection)
    coverage_database_file = None
    coverage_config_file = None
    if precomputed_behavioral is not None:
        behavioral_results = precomputed_behavioral
    else:
        (
            xml_path,
            run_result,
            coverage_database_file,
            coverage_config_file,
        ) = await async_run_behavioral_tests(
            test_files=test_files,
            test_env=test_env,
            cwd=cwd,
            pytest_cmd=test_config.pytest_cmd,
            enable_coverage=True,
        )
        behavioral_results = parse_test_results(
            test_xml_path=xml_path,
            test_files=test_files,
            test_config=test_config,
            optimization_iteration=optimization_iteration,
            run_result=run_result,
        )

    if not behavioral_results:
        if run_result is not None:
            log.warning(
                "No behavioral test results for original code (exit code %s). "
                "Skipping baseline establishment.\nSTDOUT: %.2000s\nSTDERR: %.2000s",
                run_result.returncode,
                run_result.stdout,
                run_result.stderr,
            )
        else:
            log.warning(
                "No behavioral test results for original code. "
                "Skipping baseline establishment.",
            )
        return None

    # Steps 2 & 3 run concurrently: line profiling and benchmarking
    # are independent subprocesses that don't share state.
    originals = add_async_perf_decorator(async_function, cwd)
    try:
        import asyncio  # noqa: PLC0415

        (
            (lp_xml_path, lp_run_result),
            (bm_xml_path, bm_run_result),
        ) = await asyncio.gather(
            async_run_line_profile_tests(
                test_files=test_files,
                test_env=test_env,
                cwd=cwd,
                pytest_cmd=test_config.pytest_cmd,
                result_file_name="pytest_lp_results.xml",
            ),
            async_run_benchmarking_tests(
                test_files=test_files,
                test_env=test_env,
                cwd=cwd,
                pytest_cmd=test_config.pytest_cmd,
                result_file_name="pytest_bm_results.xml",
            ),
        )
    finally:
        revert_async_decorator(originals)

    line_profile_results = parse_test_results(
        test_xml_path=lp_xml_path,
        test_files=test_files,
        test_config=test_config,
        optimization_iteration=optimization_iteration,
        run_result=lp_run_result,
    )
    benchmarking_results = parse_test_results(
        test_xml_path=bm_xml_path,
        test_files=test_files,
        test_config=test_config,
        optimization_iteration=optimization_iteration,
        run_result=bm_run_result,
    )

    # Step 4: validate benchmark runtime
    runtime_data = benchmarking_results.usable_runtime_data_by_test_case()
    total_timing = benchmarking_results.total_passed_runtime()
    if not runtime_data and not is_async:
        log.warning(
            "No benchmark timing data for original code. "
            "Skipping baseline establishment.",
        )
        return None

    # Step 5: identify failed regression tests to remove
    functions_to_remove = tuple(
        result.id.test_function_name
        for result in behavioral_results
        if result.test_type == TestType.GENERATED_REGRESSION
        and not result.did_pass
        and result.id.test_function_name is not None
    )

    loop_count = benchmarking_results.number_of_loops()
    log.info(
        "Original code runtime over %d loop%s: %d ns total",
        loop_count,
        "s" if loop_count != 1 else "",
        total_timing,
    )

    return OriginalCodeBaseline(
        behavior_test_results=behavioral_results,
        benchmarking_test_results=benchmarking_results,
        runtime=total_timing,
        line_profile_results=line_profile_results,
        functions_to_remove=functions_to_remove,
        coverage_database_file=coverage_database_file,
        coverage_config_file=coverage_config_file,
    )


def add_async_perf_decorator(
    func: FunctionToOptimize | None,
    project_root: Path,
) -> dict[Path, str]:
    """Add the async performance decorator to *func* if applicable.

    Returns the originals dict for later revert.
    """
    if func is None or not func.is_async:
        return {}

    from .._model import TestingMode  # noqa: PLC0415
    from ..testing._instrumentation import (  # noqa: PLC0415
        add_async_decorator_to_function,
    )

    added, originals = add_async_decorator_to_function(
        func.file_path,
        func,
        TestingMode.PERFORMANCE,
        project_root=project_root,
    )
    if added:
        log.info(
            "Added async performance decorator to %s",
            func.function_name,
        )
    return originals


def revert_async_decorator(originals: dict[Path, str]) -> None:
    """Revert files modified by :func:`add_async_perf_decorator`."""
    if not originals:
        return

    from ..testing._instrumentation import (  # noqa: PLC0415
        revert_instrumented_files,
    )

    revert_instrumented_files(originals)
