"""Candidate evaluation and selection for the optimization pipeline.

Contains async functions for evaluating optimization candidates
against behavioral tests and benchmarks, ranking them, and selecting
the best result.  All evaluation functions are coroutines designed
to run under a single top-level ``asyncio.run()`` call.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from codeflash_core import (
    AIClient,
    Candidate,
    EvaluationContext,
    performance_gain,
    select_best,
)

from ..codegen._replacement import replace_functions_in_file
from ..testing._parse_results import parse_test_results
from ..verification._verification import compare_test_results

if TYPE_CHECKING:
    from pathlib import Path
    from typing import Any

    from ..benchmarking.models import BenchmarkKey
    from ..testing.models import (
        TestFiles,
        TestResults,
    )
    from ..verification.models import OriginalCodeBaseline
    from ._context import OptimizationContext
    from ._optimizer import FunctionInput

log = logging.getLogger(__name__)


async def run_tests_and_benchmark(  # noqa: PLR0913
    cid: str,
    fn_input: FunctionInput,
    baseline: OriginalCodeBaseline,
    eval_ctx: EvaluationContext,
    test_files: TestFiles,
    test_env: dict[str, str],
    ctx: OptimizationContext,
    failed_candidate_diffs: dict[str, list[Any]],
    candidate_bench_results: dict[str, TestResults],
    *,
    evaluate_async_fn: _EvalAsyncFn | None = None,
    cwd: Path | None = None,
) -> float | None:
    """Run behavioral tests and benchmarks for a candidate.

    Uses ``asyncio.create_subprocess_exec`` for test execution,
    keeping all ``eval_ctx`` mutations between ``await`` points
    so they are inherently safe under cooperative concurrency.

    Expects the updated source to already be written to disk.
    When *cwd* is given, test subprocesses use that directory
    instead of ``ctx.project_root``.
    """
    from ..testing._test_runner import (  # noqa: PLC0415
        async_run_behavioral_tests,
        async_run_benchmarking_tests,
    )

    project_cwd = cwd or ctx.project_root
    rootdir = ctx.test_cfg.tests_project_rootdir

    xml_path, run_result, _, _ = await async_run_behavioral_tests(
        test_files=test_files,
        test_env=test_env,
        cwd=project_cwd,
        pytest_cmd=ctx.test_cfg.pytest_cmd,
        rootdir=rootdir,
        result_file_name=f"pytest_behavioral_{cid}.xml",
    )
    candidate_results = parse_test_results(
        test_xml_path=xml_path,
        test_files=test_files,
        test_config=ctx.test_cfg,
        optimization_iteration=0,
        run_result=run_result,
    )

    is_correct, diffs = compare_test_results(
        baseline.behavior_test_results,
        candidate_results,
    )
    if not is_correct:
        log.info(
            "Candidate %s failed behavioral tests (%d diffs)",
            cid,
            len(diffs),
        )
        if not diffs and run_result is not None and run_result.returncode != 0:
            log.warning(
                "Candidate %s test subprocess crashed"
                " (exit code %s).\nSTDOUT: %.2000s\nSTDERR:"
                " %.2000s",
                cid,
                run_result.returncode,
                run_result.stdout,
                run_result.stderr,
            )
        eval_ctx.record_failed(cid)
        if diffs:
            import attrs as _attrs  # noqa: PLC0415

            failed_candidate_diffs[cid] = [_attrs.asdict(d) for d in diffs]
        return None

    from ..verification._baseline import (  # noqa: PLC0415
        add_async_perf_decorator,
        revert_async_decorator,
    )

    func = fn_input.function
    originals = add_async_perf_decorator(
        func if func.is_async else None,
        project_cwd,
    )
    try:
        bench_xml, bench_result = await async_run_benchmarking_tests(
            test_files=test_files,
            test_env=test_env,
            cwd=project_cwd,
            pytest_cmd=ctx.test_cfg.pytest_cmd,
            rootdir=rootdir,
            result_file_name=f"pytest_bench_{cid}.xml",
        )
        bench_results = parse_test_results(
            test_xml_path=bench_xml,
            test_files=test_files,
            test_config=ctx.test_cfg,
            optimization_iteration=0,
            run_result=bench_result,
        )
    finally:
        revert_async_decorator(originals)

    optimized_runtime = bench_results.total_passed_runtime()
    is_async = fn_input.function.is_async
    candidate_bench_results[cid] = bench_results

    if not is_async and (optimized_runtime is None or optimized_runtime <= 0):
        log.debug(
            "Candidate %s has no measurable runtime",
            cid,
        )
        eval_ctx.record_failed(cid)
        return None

    if optimized_runtime is None:
        optimized_runtime = 0

    if is_async and evaluate_async_fn is not None:
        return await evaluate_async_fn(
            cid,
            fn_input,
            baseline,
            eval_ctx,
            bench_results,
            optimized_runtime,
        )

    speedup: float = performance_gain(
        original_runtime_ns=baseline.runtime,
        optimized_runtime_ns=optimized_runtime,
    )
    eval_ctx.record_success(
        cid,
        runtime=float(optimized_runtime),
        speedup=speedup,
    )
    log.info(
        "Candidate %s: %.1f%% speedup (%d ns -> %d ns)",
        cid,
        speedup * 100,
        baseline.runtime,
        optimized_runtime,
    )
    return speedup


async def evaluate_candidate_isolated(  # noqa: PLR0913
    candidate: Candidate,
    fn_input: FunctionInput,
    baseline: OriginalCodeBaseline,
    eval_ctx: EvaluationContext,
    test_files: TestFiles,
    test_env: dict[str, str],
    ctx: OptimizationContext,
    failed_candidate_code: dict[str, str],
    failed_candidate_diffs: dict[str, list[Any]],
    candidate_bench_results: dict[str, TestResults],
    *,
    evaluate_async_fn: _EvalAsyncFn | None = None,
) -> float | None:
    """Evaluate a candidate using a project overlay.

    Creates an isolated overlay directory so the original source
    is never modified.  Safe to call concurrently via
    ``asyncio.gather`` — all shared state mutations happen between
    ``await`` points (never preempted).
    """
    from ._eval_worktree import (  # noqa: PLC0415
        cleanup_overlay,
        create_project_overlay,
    )

    cid = candidate.candidate_id

    try:
        updated_source = replace_functions_in_file(
            source_code=fn_input.source_code,
            original_function_names=[
                fn_input.function.function_name,
            ],
            optimized_code=candidate.code,
            preexisting_objects=set(),
        )
    except Exception:  # noqa: BLE001
        log.info(
            "Replacement failed for candidate %s",
            cid,
            exc_info=True,
        )
        eval_ctx.record_failed(cid)
        return None

    overlay = create_project_overlay(
        fn_input.module_path,
        ctx.project_root,
        updated_source,
    )
    try:
        result = await run_tests_and_benchmark(
            cid=cid,
            fn_input=fn_input,
            baseline=baseline,
            eval_ctx=eval_ctx,
            test_files=test_files,
            test_env=test_env,
            ctx=ctx,
            failed_candidate_diffs=failed_candidate_diffs,
            candidate_bench_results=candidate_bench_results,
            evaluate_async_fn=evaluate_async_fn,
            cwd=overlay,
        )
        if result is None:
            failed_candidate_code[cid] = candidate.code
        else:
            eval_ctx.optimizations_post[cid] = candidate.code
        return result
    finally:
        cleanup_overlay(overlay)


def rank_candidates(  # noqa: PLR0913
    ai_client: AIClient,
    function_trace_id: str,
    eval_ctx: EvaluationContext,
    valid: list[Candidate],
    diff_lengths: list[int],
    original_source: str,
    original_runtime_ns: int,
) -> int | None:
    """Rank candidates, returning the index of the best one.

    Tries AI ranking first; falls back to weighted rank-sum.
    """
    import difflib  # noqa: PLC0415

    if len(valid) == 0:
        return None
    if len(valid) == 1:
        return 0

    valid_ids = [c.candidate_id for c in valid]
    original_lines = original_source.splitlines(keepends=True)
    speedups = []
    diffs = []
    for c in valid:
        sp = eval_ctx.get_speedup(c.candidate_id) or 0.0
        speedups.append(1.0 + sp)
        candidate_lines = c.code.splitlines(keepends=True)
        diff_str = "".join(
            difflib.unified_diff(original_lines, candidate_lines),
        )
        diffs.append(diff_str)

    ranking = ai_client.generate_ranking(
        trace_id=function_trace_id,
        diffs=diffs,
        candidate_ids=valid_ids,
        speedups=speedups,
    )
    if ranking:
        return ranking[0]

    best_id = select_best(
        eval_ctx,
        original_runtime_ns,
        diff_lengths,
        valid_ids,
    )
    if best_id is None:
        return None
    return next(i for i, c in enumerate(valid) if c.candidate_id == best_id)


def build_benchmark_details(  # noqa: PLR0913
    winner: Candidate,
    baseline: OriginalCodeBaseline,
    function_benchmark_timings: dict[BenchmarkKey, int],
    total_benchmark_timings: dict[BenchmarkKey, int],
    candidate_bench_results: dict[str, TestResults],
    replay_tests_dir: Path | None,
    project_root: Path,
) -> list[dict[str, object]] | None:
    """Build per-benchmark speedup details, or *None* if unavailable.

    Requires ``function_benchmark_timings``,
    ``total_benchmark_timings``, and ``replay_tests_dir``.
    Uses :meth:`TestResults.group_by_benchmarks` to compute
    per-benchmark performance gain from replay test results.
    """
    if not function_benchmark_timings or not total_benchmark_timings:
        return None

    cid = winner.candidate_id
    winner_bench = candidate_bench_results.get(cid)
    if winner_bench is None:
        return None

    from ..benchmarking._benchmarking import (  # noqa: PLC0415
        process_benchmark_data,
    )

    benchmark_keys = list(function_benchmark_timings)

    if replay_tests_dir is not None:
        orig_by_bk = baseline.benchmarking_test_results.group_by_benchmarks(
            benchmark_keys,
            replay_tests_dir,
            project_root,
        )
        opt_by_bk = winner_bench.group_by_benchmarks(
            benchmark_keys,
            replay_tests_dir,
            project_root,
        )
        replay_gain: dict[BenchmarkKey, float] = {}
        for bk in benchmark_keys:
            orig_rt = orig_by_bk[bk].total_passed_runtime()
            opt_rt = opt_by_bk[bk].total_passed_runtime()
            replay_gain[bk] = performance_gain(
                original_runtime_ns=orig_rt,
                optimized_runtime_ns=opt_rt,
            )
    else:
        orig_total = baseline.benchmarking_test_results.total_passed_runtime()
        opt_total = winner_bench.total_passed_runtime()
        if not orig_total or not opt_total:
            return None
        overall = performance_gain(
            original_runtime_ns=orig_total,
            optimized_runtime_ns=opt_total,
        )
        replay_gain = dict.fromkeys(
            benchmark_keys,
            overall,
        )

    info = process_benchmark_data(
        replay_performance_gain=replay_gain,
        fto_benchmark_timings=function_benchmark_timings,
        total_benchmark_timings=total_benchmark_timings,
    )
    if info is None:
        return None
    return [
        {
            "benchmark_name": d.benchmark_name,
            "test_function": d.test_function,
            "original_timing": d.original_timing,
            "expected_new_timing": d.expected_new_timing,
            "speedup_percent": d.speedup_percent,
        }
        for d in info.benchmark_details
    ]


def log_evaluation_results(
    ai_client: AIClient,
    function_trace_id: str,
    winner: Candidate,
    eval_ctx: EvaluationContext,
    baseline: OriginalCodeBaseline,
) -> None:
    """Log evaluation results to the AI service (fire-and-forget)."""
    from codeflash_core import (  # noqa: PLC0415
        __version__ as _core_version,
    )

    payload: dict[str, Any] = {
        "trace_id": function_trace_id,
        "speedup_ratio": eval_ctx.speedup_ratios,
        "original_runtime": baseline.runtime,
        "optimized_runtime": dict(eval_ctx.optimized_runtimes),
        "is_correct": dict(eval_ctx.is_correct),
        "optimized_line_profiler_results": dict(
            eval_ctx.line_profiler_results,
        ),
        "metadata": {
            "best_optimization_id": winner.candidate_id,
        },
        "optimizations_post": dict(eval_ctx.optimizations_post),
        "codeflash_version": _core_version,
    }
    ai_client.log_results(payload)


if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    _EvalAsyncFn = Callable[
        [
            str,  # cid
            FunctionInput,
            OriginalCodeBaseline,
            EvaluationContext,
            TestResults,
            int,  # optimized_runtime
        ],
        Awaitable[float | None],
    ]
