"""Async-specific benchmarking for the optimization pipeline.

Contains standalone functions for collecting async baseline metrics,
running concurrency benchmarks, and evaluating async candidates.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import attrs

from codeflash_core import (
    EvaluationContext,
    performance_gain,
)

from ..context.pipeline import get_code_optimization_context
from ..testing._parse_results import parse_test_results
from ..testing._test_runner import async_run_benchmarking_tests

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize
    from ..benchmarking.models import ConcurrencyMetrics
    from ..context.models import CodeOptimizationContext
    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 collect_baseline_async_metrics(  # noqa: PLR0913
    baseline: OriginalCodeBaseline,
    func: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    test_env: dict[str, str],
    test_files: TestFiles | None,
    ctx: OptimizationContext,
) -> OriginalCodeBaseline:
    """Collect async throughput and concurrency metrics.

    Returns an evolved baseline with the metrics attached.
    """
    from ..testing._stdout_parsers import (  # noqa: PLC0415
        calculate_function_throughput_from_test_results,
    )

    async_throughput = calculate_function_throughput_from_test_results(
        baseline.benchmarking_test_results,
        func.function_name,
    )
    log.info(
        "Async baseline throughput: %d calls",
        async_throughput,
    )

    concurrency_metrics = await run_concurrency_benchmark(
        func=func,
        code_context=code_context,
        test_env=test_env,
        test_files=test_files,
        ctx=ctx,
    )
    if concurrency_metrics:
        log.info(
            "Baseline concurrency: ratio=%.2f, seq=%dns, conc=%dns",
            concurrency_metrics.concurrency_ratio,
            concurrency_metrics.sequential_time_ns,
            concurrency_metrics.concurrent_time_ns,
        )
    else:
        log.info("Baseline concurrency benchmark returned no metrics")

    return attrs.evolve(
        baseline,
        async_throughput=async_throughput,
        concurrency_metrics=concurrency_metrics,
    )


async def run_concurrency_benchmark(
    func: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    test_env: dict[str, str],
    test_files: TestFiles | None,
    ctx: OptimizationContext,
) -> ConcurrencyMetrics | None:
    """Run concurrency benchmark for an async function.

    Instruments the source with a concurrency decorator,
    runs performance tests, parses the metrics, and restores
    the original source.
    """
    if not func.is_async:
        return None

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

    originals: dict[Path, str] = {}
    try:
        added, originals = add_async_decorator_to_function(
            func.file_path,
            func,
            TestingMode.CONCURRENCY,
            project_root=ctx.project_root,
        )
        if not added:
            log.info(
                "Concurrency decorator not added to %s",
                func.function_name,
            )
            return None

        if test_files is None:
            return None

        bench_xml, bench_result = await async_run_benchmarking_tests(
            test_files=test_files,
            test_env=test_env,
            cwd=ctx.project_root,
            pytest_cmd=ctx.test_cfg.pytest_cmd,
            min_loops=1,
            max_loops=3,
            target_duration_seconds=5.0,
        )
        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,
        )
    except Exception:  # noqa: BLE001
        log.info(
            "Concurrency benchmark failed",
            exc_info=True,
        )
        return None
    finally:
        if originals:
            revert_instrumented_files(originals)

    return parse_concurrency_metrics(
        bench_results,
        func.function_name,
    )


async def evaluate_async_candidate(  # noqa: PLR0913
    cid: str,
    fn_input: FunctionInput,
    baseline: OriginalCodeBaseline,
    eval_ctx: EvaluationContext,
    bench_results: TestResults,
    optimized_runtime: int,
    test_files: TestFiles | None,
    ctx: OptimizationContext,
) -> tuple[float | None, str | None]:
    """Evaluate an async candidate using throughput and concurrency.

    Returns *(speedup, acceptance_reason)*.  *speedup* is ``None``
    when the candidate is rejected.
    """
    from ..testing._stdout_parsers import (  # noqa: PLC0415
        calculate_function_throughput_from_test_results,
    )
    from ..verification._critic import (  # noqa: PLC0415
        get_acceptance_reason,
        speedup_critic,
    )
    from ..verification.models import (  # noqa: PLC0415
        OptimizedCandidateResult,
    )
    from ._test_orchestrator import build_test_env  # noqa: PLC0415

    func = fn_input.function
    candidate_throughput = calculate_function_throughput_from_test_results(
        bench_results,
        func.function_name,
    )

    candidate_concurrency = await run_concurrency_benchmark(
        func,
        get_code_optimization_context(func, ctx.project_root),
        build_test_env(fn_input, ctx.project_root, ctx.test_cfg),
        test_files,
        ctx,
    )

    candidate_result = OptimizedCandidateResult(
        max_loop_count=bench_results.number_of_loops(),
        best_test_runtime=optimized_runtime,
        behavior_test_results=bench_results,
        benchmarking_test_results=bench_results,
        optimization_candidate_index=0,
        total_candidate_timing=optimized_runtime,
        async_throughput=candidate_throughput,
        concurrency_metrics=candidate_concurrency,
    )

    log.info(
        "Async candidate %s: throughput=%d, concurrency=%s, runtime=%d",
        cid,
        candidate_throughput,
        candidate_concurrency,
        optimized_runtime,
    )

    accepted = speedup_critic(
        candidate_result,
        baseline.runtime,
        None,
        original_async_throughput=baseline.async_throughput,
        original_concurrency_metrics=(baseline.concurrency_metrics),
    )
    if not accepted:
        log.info("Candidate %s rejected by async critic", cid)
        eval_ctx.record_failed(cid)
        return None, None

    reason = get_acceptance_reason(
        baseline.runtime,
        optimized_runtime,
        original_async_throughput=baseline.async_throughput,
        optimized_async_throughput=candidate_throughput,
        original_concurrency_metrics=(baseline.concurrency_metrics),
        optimized_concurrency_metrics=candidate_concurrency,
    )
    log.info(
        "Candidate %s accepted for reason: %s",
        cid,
        reason.value,
    )

    # Use a synthetic speedup for ranking purposes.
    # For async, factor in all available dimensions.
    speedup = performance_gain(
        original_runtime_ns=max(baseline.runtime, 1),
        optimized_runtime_ns=max(optimized_runtime, 1),
    )
    if candidate_concurrency and baseline.concurrency_metrics:
        baseline_ratio = baseline.concurrency_metrics.concurrency_ratio
        speedup = max(
            speedup,
            (candidate_concurrency.concurrency_ratio - baseline_ratio)
            / max(baseline_ratio, 1.0),
        )
    if (
        baseline.async_throughput is not None
        and candidate_throughput > 0
        and baseline.async_throughput > 0
    ):
        speedup = max(
            speedup,
            (candidate_throughput - baseline.async_throughput)
            / baseline.async_throughput,
        )

    eval_ctx.record_success(
        cid,
        runtime=float(optimized_runtime),
        speedup=speedup,
    )
    eval_ctx.async_throughputs[cid] = candidate_throughput
    if candidate_concurrency is not None:
        eval_ctx.candidate_concurrency[cid] = candidate_concurrency
    log.info(
        "Candidate %s: %s improvement (%.1f%%)",
        cid,
        reason.value,
        speedup * 100,
    )
    return speedup, reason.value
