"""Per-function optimization loop for Python code.

Contains standalone helper functions (numerical detection, AST
resolution, code replacement) **and** the composable
:class:`PythonFunctionOptimizer` orchestrator that ties them
together with core pipeline building blocks.

The heavy lifting is delegated to focused modules:

* :mod:`._candidate_gen` -- candidate generation strategies
* :mod:`._candidate_eval` -- evaluation, ranking, selection
* :mod:`._test_orchestrator` -- test instrumentation & generation
* :mod:`._async_bench` -- async-specific benchmarking
"""

from __future__ import annotations

import ast
import importlib.util
import logging
from typing import TYPE_CHECKING

import attrs

from codeflash_core import (
    Candidate,
    EvaluationContext,
    OptimizationReviewResult,
    dedup_candidates,
    diff_length,
    humanize_runtime,
)
from codeflash_core import (
    __version__ as _core_version,
)

from .._constants import LANGUAGE_FIELDS, LANGUAGE_VERSION

if TYPE_CHECKING:
    from pathlib import Path
    from typing import Any

    from .._model import FunctionParent, FunctionToOptimize
    from ..benchmarking.models import BenchmarkKey
    from ..context.models import (
        CodeOptimizationContext,
        CodeStringsMarkdown,
    )
    from ..test_discovery.models import FunctionCalledInTest, TestType
    from ..testing.models import (
        TestFile,
        TestFiles,
        TestResults,
    )
    from ..verification.models import OriginalCodeBaseline
    from ._context import OptimizationContext
    from ._optimizer import FunctionInput, FunctionResult

log = logging.getLogger(__name__)

_MAX_TEST_EXAMPLES_CHARS = 4000


def _collect_test_sources(
    test_files: TestFiles,
    test_type: TestType,
    budget: int,
) -> tuple[list[str], int]:
    """Read test file sources of a given *test_type* within *budget* chars."""
    parts: list[str] = []
    used = 0
    for tf in test_files.test_files:
        if tf.test_type != test_type:
            continue
        try:
            source = tf.original_file_path.read_text(
                encoding="utf-8",
            )
        except Exception:  # noqa: BLE001, S112
            continue
        if not source.strip():
            continue
        remaining = budget - used
        if remaining <= 0:
            break
        chunk = source[:remaining]
        parts.append(chunk)
        used += len(chunk)
    return parts, used


def _extract_test_input_examples(
    test_files: TestFiles | None,
) -> str | None:
    """Extract test source code as behavioral examples for the optimizer.

    Includes both EXISTING_UNIT_TEST (hand-written assertions) and
    GENERATED_REGRESSION test files, truncated to
    *_MAX_TEST_EXAMPLES_CHARS*.  Existing tests are prioritized
    since they represent the developer's explicit behavioral
    expectations.

    Returns *None* when no test source is available.
    """
    if test_files is None:
        return None

    from ..test_discovery.models import TestType  # noqa: PLC0415

    existing, used = _collect_test_sources(
        test_files,
        TestType.EXISTING_UNIT_TEST,
        _MAX_TEST_EXAMPLES_CHARS,
    )
    generated, _ = _collect_test_sources(
        test_files,
        TestType.GENERATED_REGRESSION,
        _MAX_TEST_EXAMPLES_CHARS - used,
    )

    if not existing and not generated:
        return None

    sections: list[str] = []
    if existing:
        sections.append(
            "## Hand-written unit tests"
            " (developer expectations):\n" + "\n\n".join(existing),
        )
    if generated:
        sections.append(
            "## Generated regression tests:\n" + "\n\n".join(generated),
        )
    return "\n\n".join(sections)


_HAS_NUMBA: bool = importlib.util.find_spec("numba") is not None

NUMERICAL_MODULES: frozenset[str] = frozenset(
    {"numpy", "torch", "numba", "jax", "tensorflow", "math", "scipy"},
)

NUMBA_REQUIRED_MODULES: frozenset[str] = frozenset(
    {"numpy", "math", "scipy"},
)


async def _evaluate_batch_parallel(
    candidates: list[Candidate],
    try_fn: _AsyncTryCandidateFn,
    max_workers: int = 4,
) -> None:
    """Evaluate *candidates* concurrently using asyncio.

    Each candidate runs as a coroutine; a semaphore limits
    concurrency to *max_workers*.  Because asyncio is cooperative,
    all shared-state mutations between ``await`` points are
    inherently safe without locks.

    Must be called from within an already-running event loop
    (i.e. under a single top-level ``asyncio.run()``).
    """
    import asyncio  # noqa: PLC0415

    if len(candidates) <= 1:
        for c in candidates:
            await try_fn(c)
        return

    sem = asyncio.Semaphore(max_workers)

    async def _guarded(c: Candidate) -> None:
        async with sem:
            try:
                await try_fn(c)
            except Exception:  # noqa: BLE001
                log.warning(
                    "Candidate %s raised during evaluation",
                    c.candidate_id,
                    exc_info=True,
                )

    await asyncio.gather(
        *(_guarded(c) for c in candidates),
    )


def _uses_numerical_names(node: ast.AST, numerical_names: set[str]) -> bool:
    """Return *True* if *node* references any of *numerical_names*."""
    return any(
        isinstance(n, ast.Name) and n.id in numerical_names
        for n in ast.walk(node)
    )


def _collect_numerical_imports(
    tree: ast.Module,
) -> tuple[set[str], set[str]]:
    """Collect imported names and root modules from numerical libraries."""
    numerical_names: set[str] = set()
    modules_used: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                module_root = alias.name.split(".")[0]
                if module_root in NUMERICAL_MODULES:
                    numerical_names.add(alias.asname or module_root)
                    modules_used.add(module_root)
        elif isinstance(node, ast.ImportFrom) and node.module:
            module_root = node.module.split(".")[0]
            if module_root in NUMERICAL_MODULES:
                for alias in node.names:
                    if alias.name == "*":
                        numerical_names.add(module_root)
                    else:
                        numerical_names.add(alias.asname or alias.name)
                modules_used.add(module_root)
    return numerical_names, modules_used


def _find_function_node(
    tree: ast.Module, name_parts: list[str]
) -> ast.FunctionDef | None:
    """Find a function node by qualified name parts.

    E.g. ``["Class", "method"]``.
    """
    if not name_parts or len(name_parts) > 2:  # noqa: PLR2004
        return None
    body: list[ast.stmt] = tree.body
    for part in name_parts[:-1]:
        for node in body:
            if isinstance(node, ast.ClassDef) and node.name == part:
                body = node.body
                break
        else:
            return None
    for node in body:
        if isinstance(node, ast.FunctionDef) and node.name == name_parts[-1]:
            return node
    return None


def is_numerical_code(
    code_string: str, function_name: str | None = None
) -> bool:
    """Check if code uses numerical computing libraries.

    Detects usage of numpy, torch, numba, jax, tensorflow, scipy,
    and math.  Returns ``False`` for math/numpy/scipy when numba
    is not installed, since those optimizations require numba.
    """
    try:
        tree = ast.parse(code_string)
    except SyntaxError:
        return False

    numerical_names, modules_used = _collect_numerical_imports(tree)

    if not function_name:
        return bool(modules_used) and (
            _HAS_NUMBA or not modules_used.issubset(NUMBA_REQUIRED_MODULES)
        )

    name_parts = function_name.split(".")
    target_function = _find_function_node(tree, name_parts)
    if target_function is None:
        return False

    if not _uses_numerical_names(target_function, numerical_names):
        return False

    return not (
        not _HAS_NUMBA and modules_used.issubset(NUMBA_REQUIRED_MODULES)
    )


def resolve_function_ast(
    source_code: str,
    function_name: str,
    parents: list[FunctionParent],
) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
    """Parse *source_code* and look up a function by name and parents."""
    from ._module_prep import resolve_python_function_ast  # noqa: PLC0415

    module_ast = ast.parse(source_code)
    return resolve_python_function_ast(function_name, parents, module_ast)


def replace_function_and_helpers(  # noqa: PLR0913
    source_code: str,
    original_function_names: list[str],
    optimized_code: str,
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
    original_helper_code: dict[Path, str],
    function_to_optimize: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    optimized_code_markdown: CodeStringsMarkdown,
    project_root: Path,
) -> str:
    """Replace function definitions and revert unused helpers.

    Calls :func:`replace_functions_in_file` for the replacement,
    then detects and reverts any helper functions introduced by
    the optimizer that turned out to be unused.

    Returns the updated source code.
    """
    from ..codegen._replacement import (  # noqa: PLC0415
        replace_functions_in_file,
    )
    from ..verification._unused_helpers import (  # noqa: PLC0415
        detect_unused_helper_functions,
        revert_unused_helper_functions,
    )

    updated = replace_functions_in_file(
        source_code=source_code,
        original_function_names=original_function_names,
        optimized_code=optimized_code,
        preexisting_objects=preexisting_objects,
    )

    # Detect and revert unused helpers
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        optimized_code_markdown,
    )
    if unused_helpers:
        revert_unused_helper_functions(
            project_root, unused_helpers, original_helper_code
        )

    return updated


def apply_optimized_code(
    function_to_optimize: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    optimized_code_str: str,
    original_helper_code: dict[Path, str],
    project_root: Path,
) -> None:
    """Apply optimized code from a markdown string to files on disk.

    Groups the target function and its helpers by file, then for
    each file: adds global assignments, replaces function
    definitions, and writes back.  Finally detects and reverts
    unused helpers.
    """
    import pathlib  # noqa: PLC0415
    from collections import defaultdict  # noqa: PLC0415

    from ..codegen._replacement import (  # noqa: PLC0415
        add_global_assignments,
        is_zero_diff,
        replace_functions_and_add_imports,
    )
    from ..context.models import (  # noqa: PLC0415
        CodeStringsMarkdown,
    )

    # Build function-names-by-file, matching the original
    # group_functions_by_file logic.
    functions_by_file: dict[Path, set[str]] = defaultdict(set)
    functions_by_file[function_to_optimize.file_path].add(
        function_to_optimize.qualified_name
    )
    for helper in code_context.helper_functions:
        if helper.definition_type in ("function", None):
            functions_by_file[helper.file_path].add(
                helper.qualified_name,
            )

    markdown = CodeStringsMarkdown.parse_markdown_code(optimized_code_str)

    for module_abspath, qualified_names in functions_by_file.items():
        # Find the optimized code block for this module.
        rel = module_abspath.relative_to(project_root)
        code_to_apply: str | None = None
        for cs in markdown.code_strings:
            if cs.file_path is not None and pathlib.Path(cs.file_path) == rel:
                code_to_apply = cs.code
                break
        if code_to_apply is None:
            continue
        if not module_abspath.exists():
            continue

        source = module_abspath.read_text(encoding="utf-8")
        # Add global assignments first, then replace functions.
        source_with_globals = add_global_assignments(
            code_to_apply,
            source,
        )
        updated = replace_functions_and_add_imports(
            source_code=source_with_globals,
            function_names=list(qualified_names),
            optimized_code=code_to_apply,
            module_abspath=module_abspath,
            preexisting_objects=code_context.preexisting_objects,
            project_root_path=project_root,
        )
        if not is_zero_diff(source, updated):
            module_abspath.write_text(updated, encoding="utf-8")

    from ..verification._unused_helpers import (  # noqa: PLC0415
        detect_unused_helper_functions,
        revert_unused_helper_functions,
    )

    unused = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        markdown,
    )
    fto_qname = function_to_optimize.qualified_name
    unused = [h for h in unused if h.qualified_name != fto_qname]
    if unused:
        revert_unused_helper_functions(
            project_root,
            unused,
            original_helper_code,
        )


def _generate_trace_id() -> str:
    """Return a fresh UUID trace ID for an optimization session."""
    import uuid  # noqa: PLC0415

    return str(uuid.uuid4())


@attrs.define
class PythonFunctionOptimizer:
    """Per-function optimization loop for Python.

    Composes core pipeline building blocks
    (:class:`~codeflash_core.EvaluationContext`,
    :func:`~codeflash_core.dedup_candidates`,
    :func:`~codeflash_core.select_best`) with Python-specific
    functions (context extraction, normalization, comparison)
    into a complete per-function optimization run.

    Designed to be used as the *optimize_fn* argument to
    :meth:`PythonOptimizer.run`::

        ctx = OptimizationContext(
            plugin=plugin,
            project_root=project_root,
            test_cfg=test_cfg,
            ai_client=ai_client,
        )
        fn_opt = PythonFunctionOptimizer(ctx=ctx)
        results = project_optimizer.run(
            file_to_funcs=discovered,
            optimize_fn=fn_opt.optimize,
        )
    """

    ctx: OptimizationContext
    test_files: TestFiles | None = None
    function_to_tests: dict[str, set[FunctionCalledInTest]] | None = None
    acceptance_reason: str | None = None
    function_trace_id: str = attrs.Factory(_generate_trace_id)
    baseline_lp_markdown: str = ""
    coverage_message: str = ""
    candidate_bench_results: dict[str, TestResults] = attrs.Factory(dict)
    _function_references_cache: str | None = None
    replay_tests_dir: Path | None = None
    function_benchmark_timings: dict[BenchmarkKey, int] = attrs.Factory(dict)
    total_benchmark_timings: dict[BenchmarkKey, int] = attrs.Factory(dict)
    failed_candidate_diffs: dict[str, list[Any]] = attrs.Factory(dict)
    failed_candidate_code: dict[str, str] = attrs.Factory(dict)
    language_version: str = LANGUAGE_VERSION
    no_gen_tests: bool = False
    _concolic_dir: Path | None = None
    _last_review_tests: (
        tuple[str, str, str, str, str, int, int, int] | None
    ) = None

    async def optimize(  # noqa: C901, PLR0912, PLR0915
        self,
        fn_input: FunctionInput,
    ) -> FunctionResult:
        """Run the per-function optimization pipeline.

        1. Extract code optimization context.
        2. Detect numerical code characteristics.
        2b. Instrument tests for this function.
        3. Establish original code baseline (behavioral +
           performance).
        4. Generate optimization candidates from the AI service.
        5. Evaluate candidates: dedup, test, benchmark, rank.
        6. Select and return the best result.
        """
        from ..context.pipeline import (  # noqa: PLC0415
            get_code_optimization_context,
        )
        from ..verification._baseline import (  # noqa: PLC0415
            establish_original_code_baseline,
        )
        from ._async_bench import (  # noqa: PLC0415
            collect_baseline_async_metrics,
        )
        from ._candidate_gen import (  # noqa: PLC0415
            generate_candidates,
            generate_lp_candidates,
        )
        from ._module_prep import resolve_python_function_ast  # noqa: PLC0415
        from ._optimizer import FunctionResult  # noqa: PLC0415
        from ._test_orchestrator import (  # noqa: PLC0415
            build_test_env,
            cleanup_generated_files,
            generate_ai_tests,
            generate_concolic_tests,
            instrument_tests_for_function,
            load_and_log_coverage,
        )

        # Copy benchmark timings from the orchestrator input.
        self.function_benchmark_timings = fn_input.function_benchmark_timings
        self.total_benchmark_timings = fn_input.total_benchmark_timings

        func = fn_input.function

        # 1. Code context.
        try:
            code_context = get_code_optimization_context(
                func,
                self.ctx.project_root,
            )
        except ValueError as exc:
            return FunctionResult(
                function=func,
                module_path=fn_input.module_path,
                success=False,
                message=(f"Context extraction failed: {exc}"),
            )

        # 2. Numerical code detection.
        numerical = is_numerical_code(
            fn_input.source_code,
            func.qualified_name,
        )
        if numerical:
            log.debug(
                "%s uses numerical libraries",
                func.qualified_name,
            )

        # 2b. Generate concolic coverage tests, then instrument.
        func_ast = resolve_python_function_ast(
            func.function_name,
            list(func.parents),
            fn_input.module_ast,
        )
        fn_to_concolic: dict[str, set[FunctionCalledInTest]] = {}
        _concolic_code = ""
        if func_ast is not None:
            result = generate_concolic_tests(
                func, func_ast, self.ctx.project_root, self.ctx.test_cfg
            )
            fn_to_concolic = result[0]
            _concolic_code = result[1]
            self._concolic_dir = result[2]
        if fn_to_concolic and self.function_to_tests is not None:
            self.function_to_tests = {
                key: self.function_to_tests.get(key, set())
                | fn_to_concolic.get(key, set())
                for key in set(self.function_to_tests) | set(fn_to_concolic)
            }

        instrumented = instrument_tests_for_function(
            func,
            self.function_to_tests,
            self.ctx.project_root,
            self.ctx.test_cfg,
        )
        if instrumented is not None:
            self.test_files = instrumented

        # 2c. AI test generation (skip when --no-gen-tests).
        if self.no_gen_tests:
            generated_test_files: list[TestFile] = []
        else:
            generated_test_files = generate_ai_tests(
                ctx=self.ctx,
                function_trace_id=self.function_trace_id,
                language_version=self.language_version,
                func=func,
                code_context=code_context,
                fn_input=fn_input,
                is_numerical=numerical,
            )
        if generated_test_files:
            from ..testing.models import (  # noqa: PLC0415
                TestFiles,
            )

            if self.test_files is None:
                self.test_files = TestFiles(
                    test_files=generated_test_files,
                )
            else:
                self.test_files.test_files.extend(
                    generated_test_files,
                )
        log.info(
            "Generated %d tests for '%s'",
            len(generated_test_files),
            func.qualified_name,
        )

        try:
            # 3. Establish baseline.
            if self.test_files is None or not self.test_files.test_files:
                return FunctionResult(
                    function=func,
                    module_path=fn_input.module_path,
                    success=False,
                    message="No test files available",
                )

            test_env = build_test_env(
                fn_input, self.ctx.project_root, self.ctx.test_cfg
            )

            # 4. Establish baseline first so we can send runtime
            # data to the AI service for better-informed candidates.
            baseline = await establish_original_code_baseline(
                test_files=self.test_files,
                test_config=self.ctx.test_cfg,
                test_env=test_env,
                cwd=self.ctx.project_root,
                is_async=func.is_async,
                async_function=(func if func.is_async else None),
            )

            if baseline is None:
                return FunctionResult(
                    function=func,
                    module_path=fn_input.module_path,
                    success=False,
                    message="Baseline establishment failed",
                )

            test_input_examples = _extract_test_input_examples(
                self.test_files,
            )

            candidates = generate_candidates(
                ctx=self.ctx,
                function_trace_id=self.function_trace_id,
                fn_input=fn_input,
                code_context=code_context,
                baseline=baseline,
                test_input_examples=test_input_examples,
                is_numerical=numerical,
            )

            # 3a. Collect async metrics if function is async.
            if func.is_async:
                baseline = await collect_baseline_async_metrics(
                    baseline=baseline,
                    func=func,
                    code_context=code_context,
                    test_env=test_env,
                    test_files=self.test_files,
                    ctx=self.ctx,
                )

            # 3b. Load and log coverage data.
            self.coverage_message = load_and_log_coverage(
                baseline,
                func,
                code_context,
            )

            # 4b. Line-profiler-guided candidates.
            lp_cands, lp_md = await generate_lp_candidates(
                ctx=self.ctx,
                function_trace_id=self.function_trace_id,
                test_files=self.test_files,
                fn_input=fn_input,
                code_context=code_context,
                baseline=baseline,
                test_env=test_env,
            )
            candidates.extend(lp_cands)
            if lp_md:
                self.baseline_lp_markdown = lp_md

            if not candidates:
                return FunctionResult(
                    function=func,
                    module_path=fn_input.module_path,
                    success=False,
                    message=("No optimization candidates generated"),
                )

            # 5 & 6. Deduplicate, evaluate, refine, select.
            return await self._evaluate_and_select(
                candidates,
                fn_input,
                baseline,
                code_context,
            )
        finally:
            cleanup_generated_files(self.test_files, self._concolic_dir)
            self._concolic_dir = None

    # -- Evaluation and selection ------------------------------------

    async def _evaluate_and_select(  # noqa: C901, PLR0915
        self,
        candidates: list[Candidate],
        fn_input: FunctionInput,
        baseline: OriginalCodeBaseline,
        code_context: CodeOptimizationContext,
    ) -> FunctionResult:
        """Dedup candidates, evaluate, refine/repair, select."""
        from ..analysis._normalizer import (  # noqa: PLC0415
            normalize_python_code,
        )
        from ._candidate_eval import (  # noqa: PLC0415
            log_evaluation_results,
            rank_candidates,
        )
        from ._candidate_gen import (  # noqa: PLC0415
            generate_adaptive_candidate,
            generate_refinement_candidates,
            repair_failed_candidates,
        )
        from ._optimizer import FunctionResult  # noqa: PLC0415
        from ._test_orchestrator import build_test_env  # noqa: PLC0415

        func = fn_input.function
        normalized_original = normalize_python_code(
            fn_input.source_code,
        )
        unique = dedup_candidates(
            candidates,
            normalize_fn=normalize_python_code,
            original_normalized=normalized_original,
        )
        if not unique:
            return FunctionResult(
                function=func,
                module_path=fn_input.module_path,
                success=False,
                message=("All candidates duplicated the original"),
            )

        from ._candidate_eval import (  # noqa: PLC0415
            evaluate_candidate_isolated,
        )

        eval_ctx = EvaluationContext()
        valid: list[Candidate] = []
        diff_lengths: list[int] = []

        async_eval = self._make_async_evaluator()
        test_env = build_test_env(
            fn_input, self.ctx.project_root, self.ctx.test_cfg
        )

        async def _try_candidate(c: Candidate) -> None:
            """Evaluate *c* and append to *valid* if it improves."""
            sp = await evaluate_candidate_isolated(
                candidate=c,
                fn_input=fn_input,
                baseline=baseline,
                eval_ctx=eval_ctx,
                test_files=self.test_files,  # type: ignore[arg-type]
                test_env=test_env,
                ctx=self.ctx,
                failed_candidate_code=self.failed_candidate_code,
                failed_candidate_diffs=self.failed_candidate_diffs,
                candidate_bench_results=self.candidate_bench_results,
                evaluate_async_fn=async_eval,
            )
            if sp is not None and sp > 0:
                valid.append(c)
                diff_lengths.append(
                    diff_length(c.code, fn_input.source_code),
                )

        # Pass 1: evaluate initial candidates in parallel.
        await _evaluate_batch_parallel(unique, _try_candidate)

        # Pass 2: refinement + repair.
        pass2 = generate_refinement_candidates(
            ai_client=self.ctx.ai_client,
            function_trace_id=self.function_trace_id,
            baseline_lp_markdown=self.baseline_lp_markdown,
            valid=valid,
            eval_ctx=eval_ctx,
            fn_input=fn_input,
            baseline=baseline,
            code_context=code_context,
        )
        pass2.extend(
            repair_failed_candidates(
                ai_client=self.ctx.ai_client,
                function_trace_id=self.function_trace_id,
                failed_candidate_diffs=self.failed_candidate_diffs,
                failed_candidate_code=self.failed_candidate_code,
                fn_input=fn_input,
            ),
        )
        if pass2:
            pass2_unique = dedup_candidates(
                pass2,
                normalize_fn=normalize_python_code,
                original_normalized=normalized_original,
            )
            await _evaluate_batch_parallel(
                pass2_unique,
                _try_candidate,
            )

        # Pass 3: adaptive optimization (needs >=2 valid).
        if len(valid) >= 2:  # noqa: PLR2004
            adaptive = generate_adaptive_candidate(
                ai_client=self.ctx.ai_client,
                function_trace_id=self.function_trace_id,
                valid=valid,
                eval_ctx=eval_ctx,
                fn_input=fn_input,
            )
            await _evaluate_batch_parallel(
                adaptive,
                _try_candidate,
            )

        if not valid:
            return FunctionResult(
                function=func,
                module_path=fn_input.module_path,
                success=False,
                message=("No candidates passed validation"),
            )

        best_idx = rank_candidates(
            ai_client=self.ctx.ai_client,
            function_trace_id=self.function_trace_id,
            eval_ctx=eval_ctx,
            valid=valid,
            diff_lengths=diff_lengths,
            original_source=fn_input.source_code,
            original_runtime_ns=baseline.runtime,
        )
        if best_idx is None:
            return FunctionResult(
                function=func,
                module_path=fn_input.module_path,
                success=False,
                message="Selection found no improvement",
            )

        winner = valid[best_idx]
        speedup = eval_ctx.get_speedup(winner.candidate_id)
        if speedup is not None:
            pct = round(speedup * 100)
            log.info(
                "\u26a1\ufe0f Optimization successful!",
            )
            reason = self.acceptance_reason
            if reason and reason != "none":
                log.info(
                    "\U0001f4c8 %d%% %s improvement",
                    pct,
                    reason,
                )
            else:
                log.info(
                    "\U0001f4c8 %d%% improvement",
                    pct,
                )

        # Post-selection: explanation, review, logging.
        annotated_tests_str = self._build_annotated_tests(
            winner,
            baseline,
        )
        explanation_text = self._generate_explanation(
            winner,
            fn_input,
            baseline,
            eval_ctx,
            code_context,
            annotated_tests_str,
        )
        review_result = self._get_optimization_review(
            winner,
            fn_input,
            baseline,
            eval_ctx,
            explanation_text,
            annotated_tests_str,
        )
        log_evaluation_results(
            ai_client=self.ctx.ai_client,
            function_trace_id=self.function_trace_id,
            winner=winner,
            eval_ctx=eval_ctx,
            baseline=baseline,
        )

        # Build PR data from the review's cached intermediates.
        from ._optimizer import PrData  # noqa: PLC0415

        pr_data: PrData | None = None
        if self._last_review_tests is not None:
            (
                existing_tests_str,
                replay_tests_str,
                concolic_tests_str,
                gen_tests_str,
                speedup_pct,
                optimized_runtime_ns,
                original_runtime_ns,
                loop_count,
            ) = self._last_review_tests
            if optimized_runtime_ns > 0:
                raw = (original_runtime_ns / optimized_runtime_ns) - 1
            else:
                raw = 0.0
            speedup_x = f"{raw:,.2f}x"
            pr_data = PrData(
                function_trace_id=(self.function_trace_id),
                existing_tests_source=existing_tests_str,
                generated_tests_source=gen_tests_str,
                replay_tests=replay_tests_str,
                concolic_tests=concolic_tests_str,
                coverage_message=(self.coverage_message or ""),
                speedup_x=speedup_x,
                speedup_pct=speedup_pct,
                best_runtime_ns=optimized_runtime_ns,
                original_runtime_ns=original_runtime_ns,
                loop_count=loop_count,
            )

        msg = f"Speedup: {speedup:.1%}" if speedup is not None else "Optimized"
        return FunctionResult(
            function=func,
            module_path=fn_input.module_path,
            success=True,
            message=msg,
            best_candidate=winner,
            explanation=explanation_text,
            review=review_result,
            pr_data=pr_data,
        )

    # -- Async evaluator factory ------------------------------------

    def _make_async_evaluator(
        self,
    ) -> _EvalAsyncFn:
        """Create a closure that delegates to the async module.

        The closure captures *self* so it can read
        ``test_files``, ``ctx`` and write
        ``acceptance_reason``.
        """
        from ._async_bench import evaluate_async_candidate  # noqa: PLC0415

        async def _eval(  # noqa: PLR0913
            cid: str,
            fn_input: FunctionInput,
            baseline: OriginalCodeBaseline,
            eval_ctx: EvaluationContext,
            bench_results: TestResults,
            optimized_runtime: int,
        ) -> float | None:
            speedup, reason = await evaluate_async_candidate(
                cid=cid,
                fn_input=fn_input,
                baseline=baseline,
                eval_ctx=eval_ctx,
                bench_results=bench_results,
                optimized_runtime=optimized_runtime,
                test_files=self.test_files,
                ctx=self.ctx,
            )
            if reason is not None:
                self.acceptance_reason = reason
            return speedup

        return _eval

    # -- Post-selection helpers -------------------------------------

    def _get_function_references(
        self,
        fn_input: FunctionInput,
    ) -> str:
        """Return markdown-formatted function call-site references.

        Uses Jedi to find where the function is called across
        the project.  Caches the result on first call.
        """
        if self._function_references_cache is not None:
            return self._function_references_cache

        from pathlib import Path as _Path  # noqa: PLC0415

        from ..analysis._function_references import (  # noqa: PLC0415
            find_function_references,
            format_references_as_markdown,
        )

        tests_root = self.ctx.test_cfg.tests_root
        refs = find_function_references(
            fn_input.function,
            self.ctx.project_root,
            tests_root=(_Path(tests_root) if tests_root else None),
        )
        result = format_references_as_markdown(
            refs,
            fn_input.function.file_path,
            self.ctx.project_root,
        )
        self._function_references_cache = result
        return result

    def _build_annotated_tests(
        self,
        winner: Candidate,
        baseline: OriginalCodeBaseline,
    ) -> str:
        """Build annotated generated-test source with runtime comments."""
        from ..test_discovery.models import (  # noqa: PLC0415
            TestType,
        )
        from ..testing._testgen import (  # noqa: PLC0415
            GeneratedTests,
            GeneratedTestsList,
            add_runtime_comments_to_generated_tests,
            remove_functions_from_generated_tests,
        )

        cid = winner.candidate_id
        winner_bench = self.candidate_bench_results.get(cid)
        if self.test_files is None or winner_bench is None:
            return ""

        # Reconstruct GeneratedTestsList from on-disk test
        # files.
        gen_tests: list[GeneratedTests] = []
        for tf in self.test_files.test_files:
            if tf.test_type != TestType.GENERATED_REGRESSION:
                continue
            try:
                source = tf.original_file_path.read_text(encoding="utf-8")
            except Exception:  # noqa: BLE001
                log.debug(
                    "Cannot read test file %s",
                    tf.original_file_path,
                )
                continue
            gen_tests.append(
                GeneratedTests(
                    generated_original_test_source=source,
                    instrumented_behavior_test_source="",
                    instrumented_perf_test_source="",
                    behavior_file_path=(tf.original_file_path),
                    perf_file_path=(
                        tf.benchmarking_file_path or tf.original_file_path
                    ),
                ),
            )
        if not gen_tests:
            return ""

        gen_list = GeneratedTestsList(generated_tests=tuple(gen_tests))

        # Runtime data.
        orig_runtimes = baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
        opt_runtimes = winner_bench.usable_runtime_data_by_test_case()

        # Annotate with runtime comments.
        from pathlib import Path as _Path  # noqa: PLC0415

        tests_rootdir = self.ctx.test_cfg.tests_root
        annotated = add_runtime_comments_to_generated_tests(
            gen_list,
            orig_runtimes,
            opt_runtimes,
            tests_project_rootdir=(
                _Path(tests_rootdir) if tests_rootdir else None
            ),
        )

        # Remove failing test functions.
        if baseline.functions_to_remove:
            annotated = remove_functions_from_generated_tests(
                annotated,
                list(baseline.functions_to_remove),
            )

        # Format as markdown code blocks.
        parts: list[str] = []
        for gt in annotated.generated_tests:
            src = gt.generated_original_test_source.strip()
            if src:
                parts.append(f"```python\n{src}\n```")
        return "\n\n".join(parts)

    def _generate_explanation(  # noqa: PLR0913
        self,
        winner: Candidate,
        fn_input: FunctionInput,
        baseline: OriginalCodeBaseline,
        eval_ctx: EvaluationContext,
        code_context: CodeOptimizationContext,
        annotated_tests_str: str,
    ) -> str:
        """Request a refined explanation from the AI service.

        Returns the new explanation, or the candidate's original
        explanation as a fallback.
        """
        speedup = eval_ctx.get_speedup(winner.candidate_id)
        speedup_pct = f"{int(speedup * 100)}%" if speedup is not None else "0%"
        optimized_runtime = eval_ctx.get_runtime(winner.candidate_id)
        lp_original = self.baseline_lp_markdown or ""
        lp_optimized = eval_ctx.line_profiler_results.get(
            winner.candidate_id,
            "",
        )
        # Async throughput metrics.
        original_throughput_str: str | None = None
        optimized_throughput_str: str | None = None
        throughput_improvement_str: str | None = None
        cid = winner.candidate_id
        candidate_throughput = eval_ctx.async_throughputs.get(cid)
        if (
            fn_input.function.is_async
            and baseline.async_throughput is not None
            and candidate_throughput is not None
        ):
            from ..verification._critic import (  # noqa: PLC0415
                throughput_gain,
            )

            original_throughput_str = (
                f"{baseline.async_throughput} operations/second"
            )
            optimized_throughput_str = (
                f"{candidate_throughput} operations/second"
            )
            tp_gain = throughput_gain(
                original_throughput=baseline.async_throughput,
                optimized_throughput=candidate_throughput,
            )
            throughput_improvement_str = f"{tp_gain * 100:.1f}%"

        # Concurrency metrics.
        original_concurrency_str: str | None = None
        optimized_concurrency_str: str | None = None
        concurrency_improvement_str: str | None = None
        candidate_conc = eval_ctx.candidate_concurrency.get(cid)
        if (
            baseline.concurrency_metrics is not None
            and candidate_conc is not None
        ):
            from ..verification._critic import (  # noqa: PLC0415
                concurrency_gain,
            )

            original_concurrency_str = (
                f"{baseline.concurrency_metrics.concurrency_ratio:.2f}x"
            )
            optimized_concurrency_str = (
                f"{candidate_conc.concurrency_ratio:.2f}x"
            )
            conc_gain = concurrency_gain(
                baseline.concurrency_metrics,
                candidate_conc,
            )
            concurrency_improvement_str = f"{conc_gain * 100:.1f}%"

        payload: dict[str, Any] = {
            "trace_id": self.function_trace_id,
            "source_code": fn_input.source_code,
            "optimized_code": winner.code,
            "dependency_code": code_context.read_only,
            "original_line_profiler_results": lp_original,
            "optimized_line_profiler_results": lp_optimized,
            "original_code_runtime": humanize_runtime(
                int(baseline.runtime),
            ),
            "optimized_code_runtime": humanize_runtime(
                int(optimized_runtime or 0),
            ),
            "speedup": speedup_pct,
            "annotated_tests": annotated_tests_str,
            "optimization_id": winner.candidate_id,
            "original_explanation": winner.explanation,
            "original_throughput": original_throughput_str,
            "optimized_throughput": optimized_throughput_str,
            "throughput_improvement": (throughput_improvement_str),
            "function_references": (
                self._get_function_references(fn_input) or None
            ),
            "acceptance_reason": (self.acceptance_reason or "runtime"),
            "original_concurrency_ratio": (original_concurrency_str),
            "optimized_concurrency_ratio": (optimized_concurrency_str),
            "concurrency_improvement": (concurrency_improvement_str),
            "codeflash_version": _core_version,
            "call_sequence": 1,
        }
        new_explanation = self.ctx.ai_client.generate_explanation(payload)
        return new_explanation or winner.explanation

    def _get_optimization_review(  # noqa: PLR0913
        self,
        winner: Candidate,
        fn_input: FunctionInput,
        baseline: OriginalCodeBaseline,
        eval_ctx: EvaluationContext,
        explanation_text: str,
        annotated_tests_str: str,
    ) -> OptimizationReviewResult:
        """Request an optimization quality review from the AI service.

        Also stores ``_last_review_tests`` on *self* so that
        the caller can build :class:`PrData` without
        recomputing.
        """
        from ._candidate_eval import build_benchmark_details  # noqa: PLC0415

        speedup = eval_ctx.get_speedup(winner.candidate_id)
        speedup_pct = f"{speedup * 100:.2f}%" if speedup is not None else "0%"
        optimized_runtime = eval_ctx.get_runtime(winner.candidate_id)
        loop_count = baseline.benchmarking_test_results.number_of_loops()

        # Build existing/replay/concolic test tables.
        existing_tests_str = ""
        replay_tests_str = ""
        concolic_tests_str = ""
        winner_bench = self.candidate_bench_results.get(
            winner.candidate_id,
        )
        if self.function_to_tests and winner_bench is not None:
            from ..codegen._create_pr import (  # noqa: PLC0415
                existing_tests_source_for,
            )

            fqn = fn_input.function.qualified_name_with_modules_from_root(
                self.ctx.project_root,
            )
            orig_runtimes = baseline.benchmarking_test_results.usable_runtime_data_by_test_case()
            opt_runtimes = winner_bench.usable_runtime_data_by_test_case()
            (
                existing_tests_str,
                replay_tests_str,
                concolic_tests_str,
            ) = existing_tests_source_for(
                fqn,
                self.function_to_tests,
                self.ctx.test_cfg,
                orig_runtimes,
                opt_runtimes,
                test_files_registry=self.test_files,
            )

        # Store for PrData construction by the caller.
        self._last_review_tests = (
            existing_tests_str,
            replay_tests_str,
            concolic_tests_str,
            annotated_tests_str,
            speedup_pct,
            int(optimized_runtime or 0),
            int(baseline.runtime),
            loop_count,
        )

        payload: dict[str, Any] = {
            **LANGUAGE_FIELDS,
            "original_code": fn_input.source_code,
            "optimized_code": winner.code,
            "explanation": explanation_text,
            "existing_tests": existing_tests_str,
            "generated_tests": annotated_tests_str,
            "trace_id": self.function_trace_id,
            "coverage_message": (self.coverage_message or ""),
            "replay_tests": replay_tests_str,
            "speedup": speedup_pct,
            "loop_count": loop_count,
            "benchmark_details": build_benchmark_details(
                winner=winner,
                baseline=baseline,
                function_benchmark_timings=(self.function_benchmark_timings),
                total_benchmark_timings=(self.total_benchmark_timings),
                candidate_bench_results=(self.candidate_bench_results),
                replay_tests_dir=self.replay_tests_dir,
                project_root=self.ctx.project_root,
            ),
            "optimized_runtime": humanize_runtime(
                int(optimized_runtime or 0),
            ),
            "original_runtime": humanize_runtime(
                int(baseline.runtime),
            ),
            "calling_fn_details": (
                self._get_function_references(fn_input) or ""
            ),
            "codeflash_version": _core_version,
            "call_sequence": 1,
        }
        result = self.ctx.ai_client.get_optimization_review(payload)
        if result.review:
            log.info(
                "Optimization review: %s",
                result.review,
            )
        return result


def write_code_and_helpers(
    code: str,
    helper_code: dict[Path, str],
    file_path: Path,
) -> None:
    """Write optimised code and its helper modules back to disk."""
    file_path.write_text(code, encoding="utf-8")
    for helper_path, content in helper_code.items():
        helper_path.write_text(content, encoding="utf-8")


if TYPE_CHECKING:
    from collections.abc import Callable, Coroutine

    _AsyncTryCandidateFn = Callable[[Candidate], Coroutine[None, None, None]]

    from ._candidate_eval import _EvalAsyncFn
