"""Test instrumentation and generation for the optimization pipeline.

Contains standalone functions for instrumenting tests, generating AI
and concolic tests, reviewing/repairing generated tests, building
test environments, loading coverage data, and cleaning up generated
files.
"""

from __future__ import annotations

import importlib.util
import logging
from typing import TYPE_CHECKING

from .._constants import LANGUAGE_FIELDS
from ..test_discovery.linking import module_name_from_file_path

if TYPE_CHECKING:
    import ast
    from pathlib import Path
    from typing import Any

    from .._model import FunctionToOptimize
    from ..context.models import CodeOptimizationContext
    from ..test_discovery.models import FunctionCalledInTest
    from ..testing.models import TestConfig, TestFile, TestFiles
    from ..verification.models import OriginalCodeBaseline
    from ._context import OptimizationContext
    from ._optimizer import FunctionInput

log = logging.getLogger(__name__)

PendingTest = tuple[
    int,  # test_index
    str,  # generated_source
    "Path",  # test_path
    "Path",  # test_perf_path
]


def load_and_log_coverage(
    baseline: OriginalCodeBaseline,
    func: FunctionToOptimize,
    code_context: CodeOptimizationContext,
) -> str:
    """Load coverage data from baseline and log it.

    Returns the coverage message string (empty when unavailable).
    """
    import os  # noqa: PLC0415

    if (
        baseline.coverage_database_file is None
        or baseline.coverage_config_file is None
    ):
        return ""

    try:
        from ..analysis._coverage import (  # noqa: PLC0415
            load_coverage_from_sqlite,
        )

        coverage_data = load_coverage_from_sqlite(
            database_path=baseline.coverage_database_file,
            config_path=baseline.coverage_config_file,
            function_name=func.qualified_name,
            code_context=code_context,
            source_code_path=func.file_path,
        )
        coverage_message = (
            f"Coverage: {coverage_data.coverage:.1f}% "
            f"for {func.qualified_name}"
        )
        log.info(
            "Coverage: %.1f%% for %s",
            coverage_data.coverage,
            func.qualified_name,
        )
        if os.environ.get("CODEFLASH_END_TO_END"):
            print(coverage_data)  # noqa: T201
    except Exception:  # noqa: BLE001
        log.debug(
            "Could not load coverage data",
            exc_info=True,
        )
        return ""
    else:
        return coverage_message


def instrument_tests_for_function(
    func: FunctionToOptimize,
    function_to_tests: (dict[str, set[FunctionCalledInTest]] | None),
    project_root: Path,
    test_cfg: TestConfig,
) -> TestFiles | None:
    """Instrument test files for *func*, returning new TestFiles.

    Checks *function_to_tests* to find which test files exercise
    *func*, then creates behavior and performance instrumented
    copies.  Returns *None* if no tests are linked or if
    *function_to_tests* is not set.
    """
    if function_to_tests is None:
        return None

    from .._model import TestingMode  # noqa: PLC0415
    from ..testing._instrumentation import (  # noqa: PLC0415
        inject_profiling_into_existing_test,
    )
    from ..testing.models import (  # noqa: PLC0415
        TestFile,
        TestFiles,
    )

    func_qname = func.qualified_name_with_modules_from_root(
        project_root,
    )
    tests_for_func = function_to_tests.get(func_qname)
    if not tests_for_func:
        return None

    from pathlib import Path as _Path  # noqa: PLC0415

    tests_project_root = _Path(test_cfg.tests_project_rootdir)
    test_file_objects: list[TestFile] = []
    seen: set[Path] = set()

    for test_info in tests_for_func:
        test_file = _Path(test_info.tests_in_file.test_file)
        if test_file in seen:
            continue
        seen.add(test_file)

        positions = [test_info.position]

        ok_beh, beh_src = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=positions,
            function_to_optimize=func,
            tests_project_root=tests_project_root,
            mode=TestingMode.BEHAVIOR,
        )
        ok_perf, perf_src = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=positions,
            function_to_optimize=func,
            tests_project_root=tests_project_root,
            mode=TestingMode.PERFORMANCE,
        )

        beh_path: Path | None = test_file.parent / (
            test_file.stem + "__perfinstrumented" + test_file.suffix
        )
        perf_path: Path | None = test_file.parent / (
            test_file.stem + "__perfonlyinstrumented" + test_file.suffix
        )

        if ok_beh and beh_src is not None:
            beh_path.write_text(beh_src, encoding="utf-8")  # type: ignore[union-attr]
        else:
            beh_path = None
        if ok_perf and perf_src is not None:
            perf_path.write_text(perf_src, encoding="utf-8")  # type: ignore[union-attr]
        else:
            perf_path = None

        test_file_objects.append(
            TestFile(
                original_file_path=test_file,
                instrumented_behavior_file_path=beh_path,
                benchmarking_file_path=perf_path,
            ),
        )

    if test_file_objects:
        n_concolic = sum(
            1
            for tf in test_file_objects
            if "test_concolic_coverage" in str(tf.original_file_path)
        )
        n_unit = len(test_file_objects) - n_concolic
        log.info(
            "Discovered %d existing unit test file%s"
            ", 0 replay test files, and"
            " %d concolic coverage test file%s for %s",
            n_unit,
            "" if n_unit == 1 else "s",
            n_concolic,
            "" if n_concolic == 1 else "s",
            func.qualified_name,
        )
        log.info(
            "Instrumented %d test file(s) for %s",
            len(test_file_objects),
            func.qualified_name,
        )
        return TestFiles(test_files=test_file_objects)
    return None


def generate_concolic_tests(
    func: FunctionToOptimize,
    func_ast: ast.FunctionDef | ast.AsyncFunctionDef,
    project_root: Path,
    test_cfg: TestConfig,
) -> tuple[dict[str, set[FunctionCalledInTest]], str, Path | None]:
    """Generate concolic coverage tests using CrossHair.

    Returns *(function_to_concolic_tests, concolic_test_code,
    concolic_dir)*.  If CrossHair is unavailable or the function
    lacks typed parameters, returns empty results.
    """
    import ast as _ast  # noqa: PLC0415
    import subprocess  # noqa: PLC0415
    import tempfile  # noqa: PLC0415

    from codeflash_python._compat import (  # noqa: PLC0415
        SAFE_SYS_EXECUTABLE,
    )

    from ..analysis._static_analysis import (  # noqa: PLC0415
        has_typed_parameters,
    )
    from ..test_discovery.discovery import (  # noqa: PLC0415
        discover_unit_tests,
    )
    from ..testing._concolic import (  # noqa: PLC0415
        clean_concolic_tests,
        is_valid_concolic_test,
        make_env_with_project_root,
    )
    from ..testing.models import TestConfig  # noqa: PLC0415

    empty: tuple[dict[str, set[FunctionCalledInTest]], str, None] = (
        {},
        "",
        None,
    )

    if not importlib.util.find_spec("crosshair"):
        log.debug(
            "Skipping concolic test generation"
            " (crosshair-tool is not installed)",
        )
        return empty

    if not isinstance(func_ast, _ast.FunctionDef) or not has_typed_parameters(
        func_ast,
        list(func.parents),
    ):
        log.debug(
            "Skipping concolic tests for %s (untyped parameters)",
            func.qualified_name,
        )
        return empty

    log.info(
        "Generating concolic opcode coverage tests"
        " for the original code\u2026",
    )

    # Build the fully-qualified function path for crosshair.
    rel = (
        func.file_path.relative_to(project_root)
        .with_suffix("")
        .as_posix()
        .replace("/", ".")
    )
    fq_target = f"{rel}.{func.qualified_name}"

    env = make_env_with_project_root(project_root)
    try:
        result = subprocess.run(  # noqa: S603
            [
                SAFE_SYS_EXECUTABLE,
                "-m",
                "crosshair",
                "cover",
                "--example_output_format=pytest",
                "--per_condition_timeout=20",
                fq_target,
            ],
            capture_output=True,
            text=True,
            cwd=str(project_root),
            check=False,
            timeout=600,
            env=env,
        )
    except subprocess.TimeoutExpired:
        log.debug("CrossHair Cover test generation timed out")
        return empty

    if result.returncode != 0:
        log.debug(
            "Error running CrossHair Cover%s",
            ": " + result.stderr if result.stderr else ".",
        )
        return empty

    generated = result.stdout
    if not is_valid_concolic_test(generated, project_root=str(project_root)):
        log.debug(
            "CrossHair generated invalid test, skipping",
        )
        return empty

    concolic_code = clean_concolic_tests(generated)

    # Write to a temp dir under the tests root so discovery
    # can find it.
    tests_root = str(test_cfg.tests_root)
    concolic_dir_str = tempfile.mkdtemp(dir=tests_root)
    from pathlib import Path as _Path  # noqa: PLC0415

    concolic_dir = _Path(concolic_dir_str)

    concolic_path = concolic_dir / "test_concolic_coverage.py"
    concolic_path.write_text(concolic_code, encoding="utf-8")

    concolic_cfg = TestConfig(
        tests_root=concolic_dir,
        tests_project_rootdir=_Path(tests_root),
        project_root_path=project_root,
        test_framework=test_cfg.test_framework,
        pytest_cmd=test_cfg.pytest_cmd,
        module_root=test_cfg.module_root,
    )
    fn_to_concolic, n_concolic, _ = discover_unit_tests(
        concolic_cfg,
    )
    log.info(
        "Created %d concolic unit test case%s",
        n_concolic,
        "s" if n_concolic != 1 else "",
    )
    return fn_to_concolic, concolic_code, concolic_dir


def generate_ai_tests(  # noqa: PLR0913
    ctx: OptimizationContext,
    function_trace_id: str,
    language_version: str,
    func: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    fn_input: FunctionInput,
    is_numerical: bool,  # noqa: FBT001
) -> list[TestFile]:
    """Generate regression tests via the AI service.

    Creates test files and instruments them client-side for
    behavior capture and performance measurement.  Returns a
    list of *TestFile* objects ready to be appended to
    ``test_files``.
    """
    import tempfile  # noqa: PLC0415
    from pathlib import Path as _Path  # noqa: PLC0415

    from codeflash_core import (  # noqa: PLC0415
        AIServiceConnectionError,
        AIServiceError,
    )

    from ..testing._testgen import generate_tests  # noqa: PLC0415

    n_tests = 2  # matches original effort default
    testgen_source = code_context.testgen_context.markdown
    if not testgen_source:
        log.debug(
            "No testgen context for %s, skipping AI test generation",
            func.qualified_name,
        )
        return []

    helper_fqns = code_context.testgen_helper_fqns or [
        h.qualified_name for h in code_context.helper_functions
    ]

    dotted_module = module_name_from_file_path(
        fn_input.module_path,
        ctx.project_root,
    )
    tests_rootdir = _Path(ctx.test_cfg.tests_project_rootdir)

    tests_root = str(ctx.test_cfg.tests_root)
    gen_dir = _Path(tempfile.mkdtemp(dir=tests_root))

    # Phase 1: generate all tests in parallel.
    from concurrent.futures import (  # noqa: PLC0415
        ThreadPoolExecutor,
        as_completed,
    )

    def _generate_one(test_index: int) -> PendingTest | None:
        test_path = gen_dir / (
            f"test__{func.function_name}__unit_test_{test_index}.py"
        )
        test_perf_path = gen_dir / (
            f"test__{func.function_name}__perf_test_{test_index}.py"
        )
        try:
            result = generate_tests(
                client=ctx.ai_client,
                source_code_being_tested=testgen_source,
                function_to_optimize=func,
                helper_function_names=helper_fqns,
                module_path=dotted_module,
                test_framework=ctx.test_cfg.test_framework,
                test_timeout=15,
                trace_id=function_trace_id,
                test_index=test_index,
                test_path=test_path,
                test_perf_path=test_perf_path,
                test_module_path=module_name_from_file_path(
                    test_path,
                    tests_rootdir,
                ),
                language_version=language_version,
                is_numerical_code=is_numerical,
            )
        except (AIServiceError, AIServiceConnectionError):
            log.debug(
                "AI service error generating test %d for %s",
                test_index,
                func.qualified_name,
                exc_info=True,
            )
            return None
        except Exception:  # noqa: BLE001
            log.debug(
                "Unexpected error generating test %d for %s",
                test_index,
                func.qualified_name,
                exc_info=True,
            )
            return None
        if result is None:
            return None
        gen_src, _raw, tp, tpp = result
        return (test_index, gen_src, tp, tpp)

    pending: list[PendingTest] = []
    with ThreadPoolExecutor(max_workers=n_tests) as pool:
        futures = {pool.submit(_generate_one, i): i for i in range(n_tests)}
        for future in as_completed(futures):
            p = future.result()
            if p is not None:
                pending.append(p)

    if not pending:
        return []

    # Phase 2+3: review and repair.
    pending = review_and_repair_tests(
        ctx=ctx,
        function_trace_id=function_trace_id,
        pending=pending,
        func=func,
        testgen_source=testgen_source,
        helper_fqns=helper_fqns,
        fn_input=fn_input,
    )

    return _write_and_instrument_tests(pending, func, tests_rootdir)


def _write_and_instrument_tests(
    pending: list[PendingTest],
    func: FunctionToOptimize,
    tests_rootdir: Path,
) -> list[TestFile]:
    """Write generated tests to disk and instrument for behavior/perf capture."""
    from .._model import TestingMode  # noqa: PLC0415
    from ..test_discovery.models import TestType  # noqa: PLC0415
    from ..testing._instrumentation import (  # noqa: PLC0415
        inject_profiling_into_existing_test,
    )
    from ..testing.models import TestFile  # noqa: PLC0415

    test_file_objects: list[TestFile] = []
    for _idx, generated_source, test_path, test_perf_path in pending:
        test_path.write_text(generated_source, encoding="utf-8")

        ok_beh, beh_src = inject_profiling_into_existing_test(
            test_path=test_path,
            call_positions=[],
            function_to_optimize=func,
            tests_project_root=tests_rootdir,
            mode=TestingMode.BEHAVIOR,
        )
        ok_perf, perf_src = inject_profiling_into_existing_test(
            test_path=test_path,
            call_positions=[],
            function_to_optimize=func,
            tests_project_root=tests_rootdir,
            mode=TestingMode.PERFORMANCE,
        )

        beh_path: Path | None = test_path.parent / (
            test_path.stem + "__perfinstrumented" + test_path.suffix
        )
        perf_path: Path | None = test_perf_path

        if ok_beh and beh_src is not None:
            beh_path.write_text(beh_src, encoding="utf-8")  # type: ignore[union-attr]
        else:
            beh_path = None
        if ok_perf and perf_src is not None:
            perf_path.write_text(perf_src, encoding="utf-8")  # type: ignore[union-attr]
        else:
            perf_path = None

        test_file_objects.append(
            TestFile(
                original_file_path=test_path,
                instrumented_behavior_file_path=beh_path,
                benchmarking_file_path=perf_path,
                test_type=TestType.GENERATED_REGRESSION,
            ),
        )

    return test_file_objects


def review_and_repair_tests(  # noqa: C901, PLR0913
    ctx: OptimizationContext,
    function_trace_id: str,
    pending: list[PendingTest],
    func: FunctionToOptimize,
    testgen_source: str,
    helper_fqns: list[str],
    fn_input: FunctionInput,
) -> list[PendingTest]:
    """Review generated tests and repair any flagged issues.

    Calls the ``/testgen_review`` endpoint; for each test with
    quality issues, calls ``/testgen_repair`` and replaces the
    in-memory sources.  Returns the (potentially updated) list.
    All errors are caught -- the pipeline never crashes here.
    """
    from codeflash_core import (  # noqa: PLC0415
        AIServiceConnectionError,
        AIServiceError,
    )

    from ..testing._testgen import (  # noqa: PLC0415
        repair_generated_tests,
        review_generated_tests,
    )

    # Build review payload.
    tests_payload: list[dict[str, Any]] = [
        {"test_index": idx, "test_source": gen_src}
        for idx, gen_src, *_ in pending
    ]
    review_payload: dict[str, Any] = {
        **LANGUAGE_FIELDS,
        "tests": tests_payload,
        "function_source_code": testgen_source,
        "function_name": func.qualified_name,
        "trace_id": function_trace_id,
    }

    try:
        reviews = review_generated_tests(
            ctx.ai_client,
            review_payload,
        )
    except (AIServiceError, AIServiceConnectionError):
        log.debug(
            "AI service error reviewing tests for %s",
            func.qualified_name,
            exc_info=True,
        )
        return pending
    except Exception:  # noqa: BLE001
        log.debug(
            "Unexpected error reviewing tests for %s",
            func.qualified_name,
            exc_info=True,
        )
        return pending

    if not reviews:
        return pending

    # Build index map for quick lookup.
    idx_to_pos = {entry[0]: pos for pos, entry in enumerate(pending)}
    from pathlib import Path as _Path  # noqa: PLC0415

    tests_root = _Path(ctx.test_cfg.tests_project_rootdir)

    for review in reviews:
        functions_to_repair = review.get("functions", [])
        if not functions_to_repair:
            continue

        review_test_index = review.get("test_index")
        if not isinstance(review_test_index, int):
            continue
        pos = idx_to_pos.get(review_test_index)
        if pos is None:
            continue

        entry = pending[pos]
        (
            tidx,
            gen_src,
            test_path,
            test_perf_path,
        ) = entry

        test_module_path = module_name_from_file_path(
            test_path,
            tests_root,
        )
        repair_payload: dict[str, Any] = {
            **LANGUAGE_FIELDS,
            "test_source": gen_src,
            "functions_to_repair": functions_to_repair,
            "function_source_code": testgen_source,
            "function_to_optimize": func.to_dict(),
            "helper_function_names": helper_fqns,
            "module_path": module_name_from_file_path(
                fn_input.module_path,
                ctx.project_root,
            ),
            "test_module_path": str(test_module_path),
            "test_framework": ctx.test_cfg.test_framework,
            "test_timeout": 15,
            "trace_id": function_trace_id,
        }

        try:
            repair_result = repair_generated_tests(
                ctx.ai_client,
                repair_payload,
            )
        except (AIServiceError, AIServiceConnectionError):
            log.debug(
                "AI service error repairing test %d for %s",
                tidx,
                func.qualified_name,
                exc_info=True,
            )
            continue
        except Exception:  # noqa: BLE001
            log.debug(
                "Unexpected error repairing test %d for %s",
                tidx,
                func.qualified_name,
                exc_info=True,
            )
            continue

        if repair_result is None:
            continue

        pending[pos] = (
            tidx,
            repair_result,
            test_path,
            test_perf_path,
        )
        log.debug(
            "Repaired test %d for %s",
            tidx,
            func.qualified_name,
        )

    return pending


def build_test_env(
    fn_input: FunctionInput,
    project_root: Path,
    test_cfg: TestConfig,
) -> dict[str, str]:
    """Build the environment for test subprocesses."""
    import os  # noqa: PLC0415
    from pathlib import Path as _Path  # noqa: PLC0415

    env = dict(os.environ)
    env["CODEFLASH_MODULE_PATH"] = str(fn_input.module_path)
    env["CODEFLASH_PROJECT_ROOT"] = str(project_root)
    # Required by instrumented tests -- the plugin overrides
    # CODEFLASH_LOOP_INDEX during looping, but a default must
    # exist before the first test function body executes.
    env["CODEFLASH_TEST_ITERATION"] = "0"
    env["CODEFLASH_LOOP_INDEX"] = "1"
    env["CODEFLASH_TRACER_DISABLE"] = "1"
    if test_cfg.module_root is not None:
        module_root = str(_Path(test_cfg.module_root))
        existing = env.get("PYTHONPATH", "")
        env["PYTHONPATH"] = (
            f"{module_root}{os.pathsep}{existing}" if existing else module_root
        )
    return env


def cleanup_generated_files(
    test_files: TestFiles | None,
    concolic_dir: Path | None,
) -> None:
    """Remove instrumented and AI-generated test files."""
    from ._orchestrator import cleanup_paths  # noqa: PLC0415

    # Always clean up the concolic temp dir, even if
    # test_files is empty (concolic dir is created before
    # instrumentation).
    if concolic_dir is not None:
        cleanup_paths([concolic_dir])

    if test_files is None:
        return
    from ..test_discovery.models import TestType  # noqa: PLC0415

    paths: list[Path | None] = []
    dirs_to_remove: set[Path] = set()
    for tf in test_files.test_files:
        paths.append(tf.instrumented_behavior_file_path)
        paths.append(tf.benchmarking_file_path)
        # Also remove original source for AI-generated tests.
        if tf.test_type == TestType.GENERATED_REGRESSION:
            paths.append(tf.original_file_path)
            dirs_to_remove.add(tf.original_file_path.parent)
    cleanup_paths(paths)
    # Remove empty temp directories created for generated tests.
    import shutil  # noqa: PLC0415

    for d in dirs_to_remove:
        shutil.rmtree(d, ignore_errors=True)
