"""Tests for post-selection orchestrator methods on PythonFunctionOptimizer."""

from __future__ import annotations

from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch

import attrs
import pytest

from codeflash_core import (
    AIClient,
    Candidate,
    EvaluationContext,
    OptimizationReviewResult,
)
from codeflash_python.pipeline._candidate_eval import (
    build_benchmark_details,
    log_evaluation_results,
)
from codeflash_python.pipeline._context import OptimizationContext
from codeflash_python.pipeline._function_optimizer import (
    PythonFunctionOptimizer,
)
from codeflash_python.testing.models import TestResults
from codeflash_python.verification.models import OriginalCodeBaseline


def _make_test_results(
    *,
    runtime: int = 100_000,
    loops: int = 1,
) -> MagicMock:
    """Build a mock TestResults with a given total runtime."""
    tr = MagicMock(spec=TestResults)
    tr.total_passed_runtime.return_value = runtime
    tr.number_of_loops.return_value = loops
    tr.usable_runtime_data_by_test_case.return_value = {}
    return tr


def _make_baseline(
    *,
    runtime: int = 200_000,
    bench_loops: int = 5,
) -> OriginalCodeBaseline:
    """Build a minimal OriginalCodeBaseline for testing."""
    behavior = _make_test_results(runtime=runtime)
    benchmarking = _make_test_results(runtime=runtime, loops=bench_loops)
    lp = _make_test_results()
    return OriginalCodeBaseline(
        behavior_test_results=behavior,
        benchmarking_test_results=benchmarking,
        runtime=runtime,
        line_profile_results=lp,
    )


def _make_candidate(
    *,
    cid: str = "cand-1",
    code: str = "def f(): return 1",
    explanation: str = "Made it faster",
) -> Candidate:
    """Build a minimal Candidate."""
    return Candidate(
        code=code,
        explanation=explanation,
        candidate_id=cid,
    )


def _make_eval_ctx(
    *,
    cid: str = "cand-1",
    speedup: float = 1.0,
    runtime: float = 100_000,
) -> EvaluationContext:
    """Build an EvaluationContext with one passing candidate recorded."""
    ctx = EvaluationContext()
    ctx.record_success(cid, runtime, speedup)
    return ctx


def _make_fn_input(
    *,
    source_code: str = "def f(): return 0",
    is_async: bool = False,
) -> MagicMock:
    """Build a mock FunctionInput."""
    fn_input = MagicMock()
    fn_input.source_code = source_code
    fn_input.module_path = Path("/tmp/module.py")
    fn_input.function.function_name = "f"
    fn_input.function.class_name = None
    fn_input.function.is_async = is_async
    fn_input.function.file_path = Path("/tmp/module.py")
    fn_input.function.qualified_name_with_modules_from_root.return_value = (
        "module.f"
    )
    return fn_input


def _make_optimizer(**overrides: Any) -> PythonFunctionOptimizer:
    """Build a PythonFunctionOptimizer with sensible test defaults."""
    ctx_defaults: dict[str, Any] = {
        "plugin": MagicMock(),
        "project_root": Path("/tmp/project"),
        "test_cfg": MagicMock(
            tests_root=None,
            module_root=None,
        ),
        "ai_client": MagicMock(spec=AIClient),
    }
    # Allow overriding ctx-level fields.
    ctx_keys = {"plugin", "project_root", "test_cfg", "ai_client"}
    ctx_overrides = {k: overrides.pop(k) for k in ctx_keys if k in overrides}
    ctx_defaults.update(ctx_overrides)
    ctx = OptimizationContext(**ctx_defaults)
    return PythonFunctionOptimizer(ctx=ctx, **overrides)


class TestGenerateExplanation:
    """Tests for _generate_explanation."""

    def test_returns_ai_explanation(self) -> None:
        """AI service explanation is returned when available."""
        opt = _make_optimizer()
        opt.ctx.ai_client.generate_explanation.return_value = (
            "Better explanation"
        )

        winner = _make_candidate()
        baseline = _make_baseline()
        eval_ctx = _make_eval_ctx()
        code_context = MagicMock(read_only="")

        result = opt._generate_explanation(
            winner,
            _make_fn_input(),
            baseline,
            eval_ctx,
            code_context,
            "",
        )

        assert "Better explanation" == result
        opt.ctx.ai_client.generate_explanation.assert_called_once()

    def test_falls_back_to_original(self) -> None:
        """Falls back to candidate explanation when AI returns empty."""
        opt = _make_optimizer()
        opt.ctx.ai_client.generate_explanation.return_value = ""

        winner = _make_candidate(explanation="Original expl")
        baseline = _make_baseline()
        eval_ctx = _make_eval_ctx()
        code_context = MagicMock(read_only="")

        result = opt._generate_explanation(
            winner,
            _make_fn_input(),
            baseline,
            eval_ctx,
            code_context,
            "",
        )

        assert "Original expl" == result

    def test_payload_contains_key_fields(self) -> None:
        """Payload sent to AI service includes required fields."""
        opt = _make_optimizer()
        opt.ctx.ai_client.generate_explanation.return_value = "ok"

        winner = _make_candidate()
        baseline = _make_baseline()
        eval_ctx = _make_eval_ctx()
        code_context = MagicMock(read_only="dep code")

        opt._generate_explanation(
            winner,
            _make_fn_input(source_code="def f(): pass"),
            baseline,
            eval_ctx,
            code_context,
            "annotated tests here",
        )

        payload = opt.ctx.ai_client.generate_explanation.call_args[0][0]
        assert "def f(): pass" == payload["source_code"]
        assert payload["optimized_code"] == winner.code
        assert "dep code" == payload["dependency_code"]
        assert "annotated tests here" == payload["annotated_tests"]
        assert payload["trace_id"] == opt.function_trace_id

    def test_async_throughput_fields(self) -> None:
        """Async throughput fields are populated for async functions."""
        opt = _make_optimizer()
        opt.ctx.ai_client.generate_explanation.return_value = "ok"

        winner = _make_candidate()
        baseline = _make_baseline()
        baseline = attrs.evolve(baseline, async_throughput=100)
        eval_ctx = _make_eval_ctx()
        eval_ctx.async_throughputs["cand-1"] = 200

        opt._generate_explanation(
            winner,
            _make_fn_input(is_async=True),
            baseline,
            eval_ctx,
            MagicMock(read_only=""),
            "",
        )

        payload = opt.ctx.ai_client.generate_explanation.call_args[0][0]
        assert "100 operations/second" == payload["original_throughput"]
        assert "200 operations/second" == payload["optimized_throughput"]
        assert payload["throughput_improvement"] is not None


class TestGetOptimizationReview:
    """Tests for _get_optimization_review."""

    def test_returns_review_result(self) -> None:
        """Review result from AI service is returned."""
        review = OptimizationReviewResult(
            review="Looks good",
            explanation="Updated",
        )
        opt = _make_optimizer()
        opt.ctx.ai_client.get_optimization_review.return_value = review

        result = opt._get_optimization_review(
            _make_candidate(),
            _make_fn_input(),
            _make_baseline(),
            _make_eval_ctx(),
            "explanation text",
            "",
        )

        assert result is review

    def test_payload_fields(self) -> None:
        """Payload includes key review fields."""
        review = OptimizationReviewResult(review="", explanation="")
        opt = _make_optimizer()
        opt.ctx.ai_client.get_optimization_review.return_value = review

        winner = _make_candidate(code="def f(): return 2")

        opt._get_optimization_review(
            winner,
            _make_fn_input(source_code="def f(): return 0"),
            _make_baseline(),
            _make_eval_ctx(),
            "my explanation",
            "generated tests md",
        )

        payload = opt.ctx.ai_client.get_optimization_review.call_args[0][0]
        assert "def f(): return 0" == payload["original_code"]
        assert "def f(): return 2" == payload["optimized_code"]
        assert "my explanation" == payload["explanation"]
        assert "generated tests md" == payload["generated_tests"]
        assert payload["speedup"] is not None
        assert payload["loop_count"] is not None
        assert "python" == payload["language"]

    def test_coverage_message_included(self) -> None:
        """Coverage message is passed through to payload."""
        review = OptimizationReviewResult(review="", explanation="")
        opt = _make_optimizer()
        opt.ctx.ai_client.get_optimization_review.return_value = review
        opt.coverage_message = "Coverage: 85.0% for module.f"

        opt._get_optimization_review(
            _make_candidate(),
            _make_fn_input(),
            _make_baseline(),
            _make_eval_ctx(),
            "",
            "",
        )

        payload = opt.ctx.ai_client.get_optimization_review.call_args[0][0]
        assert "Coverage: 85.0% for module.f" == payload["coverage_message"]


class TestLogEvaluationResults:
    """Tests for _log_evaluation_results."""

    def test_calls_log_results(self) -> None:
        """AI client log_results is called with correct payload."""
        opt = _make_optimizer()

        winner = _make_candidate()
        eval_ctx = _make_eval_ctx()
        eval_ctx.optimizations_post["cand-1"] = "def f(): return 1"
        baseline = _make_baseline()

        log_evaluation_results(
            opt.ctx.ai_client,
            opt.function_trace_id,
            winner,
            eval_ctx,
            baseline,
        )

        opt.ctx.ai_client.log_results.assert_called_once()
        payload = opt.ctx.ai_client.log_results.call_args[0][0]
        assert payload["trace_id"] == opt.function_trace_id
        assert payload["original_runtime"] == baseline.runtime
        assert "cand-1" == payload["metadata"]["best_optimization_id"]
        assert {"cand-1": "def f(): return 1"} == payload["optimizations_post"]

    def test_includes_speedup_and_correctness(self) -> None:
        """Payload includes speedup ratios and correctness flags."""
        opt = _make_optimizer()

        winner = _make_candidate()
        eval_ctx = _make_eval_ctx()
        baseline = _make_baseline()

        log_evaluation_results(
            opt.ctx.ai_client,
            opt.function_trace_id,
            winner,
            eval_ctx,
            baseline,
        )

        payload = opt.ctx.ai_client.log_results.call_args[0][0]
        assert "cand-1" in payload["speedup_ratio"]
        assert payload["is_correct"]["cand-1"] is True


class TestBuildBenchmarkDetails:
    """Tests for _build_benchmark_details."""

    def test_returns_none_without_timings(self) -> None:
        """Returns None when benchmark timings are not populated."""
        winner = _make_candidate()
        baseline = _make_baseline()

        result = build_benchmark_details(
            winner,
            baseline,
            {},
            {},
            {},
            None,
            Path("/tmp"),
        )

        assert result is None

    def test_returns_none_without_bench_results(self) -> None:
        """Returns None when candidate bench results are missing."""
        from codeflash_python.benchmarking.models import BenchmarkKey

        bk = BenchmarkKey(module_path="test_mod", function_name="test_fn")

        winner = _make_candidate()
        baseline = _make_baseline()

        result = build_benchmark_details(
            winner,
            baseline,
            {bk: 50_000},
            {bk: 200_000},
            {},
            None,
            Path("/tmp"),
        )

        assert result is None

    def test_returns_details_with_full_data(self) -> None:
        """Returns benchmark details when all data is available."""
        from codeflash_python.benchmarking.models import BenchmarkKey

        bk = BenchmarkKey(module_path="test_mod", function_name="test_fn")

        winner = _make_candidate()
        candidate_bench_results = {
            "cand-1": _make_test_results(runtime=50_000),
        }
        baseline = _make_baseline(runtime=200_000)

        result = build_benchmark_details(
            winner,
            baseline,
            {bk: 50_000},
            {bk: 200_000},
            candidate_bench_results,
            None,
            Path("/tmp"),
        )

        assert result is not None
        assert len(result) == 1
        assert "test_mod" == result[0]["benchmark_name"]
        assert "test_fn" == result[0]["test_function"]
        assert "original_timing" in result[0]
        assert "expected_new_timing" in result[0]
        assert "speedup_percent" in result[0]
