from __future__ import annotations

import uuid
from unittest.mock import AsyncMock, MagicMock

import pytest

from codeflash_api.diff._base import DiffMethod
from codeflash_api.languages.python._markdown import (
    extract_all_code_from_markdown,
    extract_code_block,
    extract_code_block_with_context,
    group_code,
    is_multi_context,
    split_markdown_code,
    truncate_pathological_output,
    wrap_code_in_markdown,
)
from codeflash_api.llm._client import LLMOutputUnparseableError
from codeflash_api.llm._models import (
    ANTHROPIC_CLAUDE_SONNET_4_5,
    OPENAI_GPT_5_MINI,
)
from codeflash_api.optimize._context import (
    BaseOptimizerContext,
    CodeStrAndExplanation,
    MultiOptimizerContext,
    SingleOptimizerContext,
    _humanize_ns,
    parse_python_version,
    validate_trace_id,
)
from codeflash_api.optimize._pipeline import (
    MAX_OPTIMIZER_CALLS,
    generate_optimization_candidate,
    get_model_distribution,
    optimize_python_code,
)
from codeflash_api.optimize.schemas import (
    OptimizedCandidateSource,
    OptimizeErrorResponse,
    OptimizeRequest,
    OptimizeResponse,
    OptimizeResponseItem,
)

from .conftest import _stub_llm_response, make_trace_id


class TestTruncatePathologicalOutput:
    """Tests for truncate_pathological_output."""

    def test_normal_code(self) -> None:
        """
        Normal code is returned unchanged.
        """
        code = "x = 1\ny = 2\n"
        assert code == truncate_pathological_output(code)

    def test_pathological(self) -> None:
        """
        Repeated escape sequences are truncated.
        """
        code = "x = 1\n\\''''''''''''more"
        result = truncate_pathological_output(code)
        assert "x = 1" in result
        assert "more" not in result


class TestExtractAllCodeFromMarkdown:
    """Tests for extract_all_code_from_markdown."""

    def test_single_block(self) -> None:
        """
        Single code block is extracted.
        """
        md = "```python\nx = 1\n```"
        assert "x = 1\n" == extract_all_code_from_markdown(md)

    def test_multiple_blocks(self) -> None:
        """
        Multiple code blocks are joined.
        """
        md = "```python\nx = 1\n```\ntext\n```python\ny = 2\n```"
        result = extract_all_code_from_markdown(md)
        assert "x = 1" in result
        assert "y = 2" in result

    def test_no_blocks(self) -> None:
        """
        No code blocks returns empty string.
        """
        assert "" == extract_all_code_from_markdown("no code here")


class TestExtractCodeBlock:
    """Tests for extract_code_block."""

    def test_basic(self) -> None:
        """
        First code block is extracted.
        """
        md = "```python\nx = 1\n```"
        assert "x = 1" == extract_code_block(md)

    def test_with_filepath(self) -> None:
        """
        Code block with filepath is extracted.
        """
        md = "```python:main.py\nx = 1\n```"
        assert "x = 1" == extract_code_block(md)

    def test_no_block(self) -> None:
        """
        No code block returns None.
        """
        assert extract_code_block("no code") is None


class TestSplitMarkdownCode:
    """Tests for split_markdown_code."""

    def test_basic(self) -> None:
        """
        Filepath blocks are parsed into a dict.
        """
        md = "```python:main.py\nx = 1\n```\n```python:util.py\ny = 2\n```"
        result = split_markdown_code(md)
        assert "main.py" in result
        assert "util.py" in result
        assert "x = 1" == result["main.py"]

    def test_plain_blocks_ignored(self) -> None:
        """
        Plain python blocks without filepath are ignored.
        """
        md = "```python\nx = 1\n```"
        result = split_markdown_code(md)
        assert {} == result

    def test_duplicate_paths(self) -> None:
        """
        First occurrence wins for duplicate paths.
        """
        md = "```python:a.py\nfirst\n```\n```python:a.py\nsecond\n```"
        result = split_markdown_code(md)
        assert "first" == result["a.py"]


class TestExtractCodeBlockWithContext:
    """Tests for extract_code_block_with_context."""

    def test_basic(self) -> None:
        """
        Code block and context are extracted.
        """
        text = "explanation\n```python\nx = 1\n```\nafter"
        result = extract_code_block_with_context(text)
        assert result is not None
        explanation, code, after = result
        assert "explanation" == explanation
        assert "x = 1\n" == code
        assert "after" == after

    def test_with_filepath(self) -> None:
        """
        Filepath variant is preferred.
        """
        text = "before\n```python:file.py\ny = 2\n```\nend"
        result = extract_code_block_with_context(text)
        assert result is not None
        assert "y = 2\n" == result[1]

    def test_no_block(self) -> None:
        """
        No code block returns None.
        """
        assert extract_code_block_with_context("no code") is None


class TestWrapCodeInMarkdown:
    """Tests for wrap_code_in_markdown."""

    def test_basic(self) -> None:
        """
        Code is wrapped in a markdown block.
        """
        assert "```python\nx = 1\n```" == wrap_code_in_markdown("x = 1")


class TestGroupCode:
    """Tests for group_code."""

    def test_single_file(self) -> None:
        """
        Single file is formatted as markdown block.
        """
        result = group_code({"main.py": "x = 1\n"})
        assert "```python:main.py" in result
        assert "x = 1" in result

    def test_multiple_files(self) -> None:
        """
        Multiple files are joined.
        """
        result = group_code({"a.py": "x = 1\n", "b.py": "y = 2\n"})
        assert "a.py" in result
        assert "b.py" in result


class TestIsMultiContext:
    """Tests for is_multi_context."""

    def test_multi(self) -> None:
        """
        Markdown with filepath header is multi.
        """
        assert is_multi_context("```python:file.py\nx=1\n```")

    def test_single(self) -> None:
        """
        Plain code is not multi.
        """
        assert not is_multi_context("x = 1\n")


class TestHumanizeNs:
    """Tests for _humanize_ns."""

    def test_nanoseconds(self) -> None:
        """
        Sub-microsecond values show ns.
        """
        assert "500ns" == _humanize_ns(500)

    def test_microseconds(self) -> None:
        """
        Microsecond range shows us.
        """
        assert "1.5us" == _humanize_ns(1500)

    def test_milliseconds(self) -> None:
        """
        Millisecond range shows ms.
        """
        assert "2.5ms" == _humanize_ns(2_500_000)

    def test_seconds(self) -> None:
        """
        Second range shows s.
        """
        assert "1.00s" == _humanize_ns(1_000_000_000)


class TestParsePythonVersion:
    """Tests for parse_python_version."""

    def test_valid(self) -> None:
        """
        Valid version is parsed to a tuple.
        """
        assert (3, 12, 9) == parse_python_version("3.12.9")

    def test_empty(self) -> None:
        """
        Empty version raises ValueError.
        """
        with pytest.raises(ValueError, match="required"):
            parse_python_version("")

    def test_wrong_format(self) -> None:
        """
        Missing parts raises ValueError.
        """
        with pytest.raises(ValueError, match="format"):
            parse_python_version("3.12")

    def test_too_old(self) -> None:
        """
        Python 3.8 is rejected.
        """
        with pytest.raises(ValueError, match="3.9"):
            parse_python_version("3.8.0")

    def test_python2(self) -> None:
        """
        Python 2 is rejected.
        """
        with pytest.raises(ValueError, match="Python 3"):
            parse_python_version("2.7.18")


class TestValidateTraceId:
    """Tests for validate_trace_id."""

    def test_valid(self) -> None:
        """
        Valid UUIDv4 passes.
        """
        valid_id = str(uuid.uuid4())
        assert validate_trace_id(valid_id)

    def test_invalid(self) -> None:
        """
        Invalid string fails.
        """
        assert not validate_trace_id("not-a-uuid")

    def test_exp_suffix(self) -> None:
        """
        EXP0 suffix is normalized and accepted.
        """
        base = str(uuid.uuid4())
        exp_id = base[:-4] + "EXP0"
        assert validate_trace_id(exp_id)


class TestCodeStrAndExplanation:
    """Tests for CodeStrAndExplanation."""

    def test_attributes(self) -> None:
        """
        Code and explanation are stored.
        """
        obj = CodeStrAndExplanation(code="x = 1", explanation="optimized")
        assert "x = 1" == obj.code
        assert "optimized" == obj.explanation


class TestBaseOptimizerContext:
    """Tests for BaseOptimizerContext."""

    def test_get_dynamic_context_single(self) -> None:
        """
        Plain source code returns SingleOptimizerContext.
        """
        ctx = BaseOptimizerContext.get_dynamic_context(
            "system", "user", "x = 1\n"
        )
        assert isinstance(ctx, SingleOptimizerContext)

    def test_get_dynamic_context_multi(self) -> None:
        """
        Multi-file markdown returns MultiOptimizerContext.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = BaseOptimizerContext.get_dynamic_context(
            "system", "user", source
        )
        assert isinstance(ctx, MultiOptimizerContext)

    def test_get_dynamic_context_single_file_multi_format(
        self,
    ) -> None:
        """
        Single file in multi format with NO_DIFF returns Single.
        """
        source = "```python:a.py\nx = 1\n```"
        ctx = BaseOptimizerContext.get_dynamic_context(
            "system", "user", source, DiffMethod.NO_DIFF
        )
        assert isinstance(ctx, SingleOptimizerContext)

    def test_is_valid_code_empty(self) -> None:
        """
        None extracted code is invalid.
        """
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        assert not ctx.is_valid_code()

    def test_validate_and_parse_valid(self) -> None:
        """
        Valid Python passes validation.
        """
        ctx = BaseOptimizerContext("s", "u", "x = 1\n")
        ctx.validate_and_parse_source_code("x = 1\n", feature_version=(3, 12))

    def test_validate_and_parse_empty(self) -> None:
        """
        Empty body raises SyntaxError.
        """
        ctx = BaseOptimizerContext("s", "u", "# comment\n")
        with pytest.raises(SyntaxError, match="empty"):
            ctx.validate_and_parse_source_code(
                "# comment\n", feature_version=(3, 12)
            )


class TestSingleOptimizerContext:
    """Tests for SingleOptimizerContext."""

    def test_get_system_prompt(self) -> None:
        """
        System prompt includes version and instructions.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "user",
            "x = 1\n",
        )
        result = ctx.get_system_prompt("3.12.9")
        assert "3.12.9" in result
        assert "FULL" in result

    def test_extract_code_from_llm_response(self) -> None:
        """
        Code block is extracted from LLM markdown response.
        """
        ctx = SingleOptimizerContext("s", "u", "x = 1\n")
        content = "explanation\n```python\ny = 2\n```"
        result = ctx.extract_code_and_explanation_from_llm_res(content)
        assert "y = 2\n" == result.code
        assert "explanation" == result.explanation

    def test_parse_generates_response_item(self) -> None:
        """
        Valid extracted code produces a response item.
        """
        ctx = SingleOptimizerContext("s", "u", "x = 1\n")
        ctx.extract_code_and_explanation_from_llm_res(
            "faster\n```python\ny = 2\n```"
        )
        result = ctx.parse_and_generate_candidate_schema()
        assert result is not None
        assert result.source_code.strip() == "y = 2"

    def test_parse_returns_none_for_empty(self) -> None:
        """
        Empty extracted code returns None.
        """
        ctx = SingleOptimizerContext("s", "u", "x = 1\n")
        ctx.extracted_code_and_expl = CodeStrAndExplanation(
            code="", explanation=""
        )
        assert ctx.parse_and_generate_candidate_schema() is None


class TestMultiOptimizerContext:
    """Tests for MultiOptimizerContext."""

    def test_get_system_prompt_no_diff(self) -> None:
        """
        NO_DIFF uses markdown format prompt.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "user",
            source,
            diff_method=DiffMethod.NO_DIFF,
        )
        result = ctx.get_system_prompt("3.12.9")
        assert "3.12.9" in result
        assert "Code Input format" in result

    def test_is_valid_code_rejects_new_files(self) -> None:
        """
        Extracted files not in original are rejected.
        """
        source = "```python:a.py\nx = 1\n```"
        ctx = MultiOptimizerContext(
            "s",
            "u",
            source,
            diff_method=DiffMethod.NO_DIFF,
        )
        ctx.extracted_code_and_expl = CodeStrAndExplanation(
            code="```python:new_file.py\nz = 3\n```",
            explanation="added",
        )
        assert not ctx.is_valid_code()


class TestGetModelDistribution:
    """Tests for get_model_distribution."""

    def test_default_five(self) -> None:
        """
        5 candidates: 3 GPT + 2 Claude.
        """
        dist = get_model_distribution(5, MAX_OPTIMIZER_CALLS)
        assert 2 == len(dist)
        gpt_model, gpt_calls = dist[0]
        claude_model, claude_calls = dist[1]
        assert 3 == gpt_calls
        assert 2 == claude_calls

    def test_capped(self) -> None:
        """
        Candidates are capped at max_calls.
        """
        dist = get_model_distribution(100, MAX_OPTIMIZER_CALLS)
        total = sum(calls for _, calls in dist)
        assert MAX_OPTIMIZER_CALLS == total

    def test_single(self) -> None:
        """
        Single candidate: 1 GPT + 0 Claude.
        """
        dist = get_model_distribution(1, MAX_OPTIMIZER_CALLS)
        total = sum(calls for _, calls in dist)
        assert 1 == total


class TestOptimizeSchemas:
    """Tests for Pydantic optimize schemas."""

    def test_request_defaults(self) -> None:
        """
        Default fields are set correctly.
        """
        req = OptimizeRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
        )
        assert 5 == req.n_candidates
        assert "python" == req.language
        assert req.is_async is False

    def test_request_version_resolution(self) -> None:
        """
        language_version backfills python_version.
        """
        req = OptimizeRequest(
            source_code="x = 1",
            trace_id=str(uuid.uuid4()),
            language_version="3.12.9",
        )
        assert "3.12.9" == req.python_version

    def test_response_item(self) -> None:
        """
        Response item is serializable.
        """
        item = OptimizeResponseItem(
            source_code="x = 2",
            explanation="faster",
            optimization_id="abc",
        )
        data = item.model_dump()
        assert "x = 2" == data["source_code"]
        assert data["parent_id"] is None

    def test_response(self) -> None:
        """
        Response wraps items.
        """
        resp = OptimizeResponse(
            optimizations=[
                OptimizeResponseItem(
                    source_code="x = 2",
                    explanation="faster",
                    optimization_id="abc",
                )
            ]
        )
        assert 1 == len(resp.optimizations)

    def test_error_response(self) -> None:
        """
        Error response has error field.
        """
        err = OptimizeErrorResponse(error="bad")
        assert "bad" == err.error

    def test_candidate_source_enum(self) -> None:
        """
        Enum values are accessible.
        """
        assert "OPTIMIZE" == OptimizedCandidateSource.OPTIMIZE.value


# ── _pipeline.py tests ──────────────────────────────────────────────


@pytest.mark.asyncio
class TestGenerateOptimizationCandidate:
    """Tests for generate_optimization_candidate."""

    async def test_success_returns_item(self) -> None:
        """
        Valid LLM response that parses to code returns an OptimizeResponseItem.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            return_value=_stub_llm_response("faster\n```python\ny = 2\n```"),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        result, cost, model_name = await generate_optimization_candidate(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
            optimize_model=OPENAI_GPT_5_MINI,
        )
        assert result is not None
        assert "y = 2" in result.source_code
        assert 0.001 == cost
        assert OPENAI_GPT_5_MINI.name == model_name

    async def test_unparseable_error_returns_none_with_cost(self) -> None:
        """
        LLMOutputUnparseableError returns None result but preserves cost.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            side_effect=LLMOutputUnparseableError(cost=0.05),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        result, cost, model_name = await generate_optimization_candidate(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
        )
        assert result is None
        assert 0.05 == cost

    async def test_generic_exception_returns_none_none_cost(self) -> None:
        """
        Unexpected exception returns None result and None cost.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            side_effect=RuntimeError("boom"),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        result, cost, model_name = await generate_optimization_candidate(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
        )
        assert result is None
        assert cost is None

    async def test_invalid_code_response_returns_none(self) -> None:
        """
        LLM returns content with no code block, result is None.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            return_value=_stub_llm_response("no code block here"),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        result, cost, model_name = await generate_optimization_candidate(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
        )
        assert result is None
        assert cost is not None

    async def test_syntax_error_in_response_returns_none(self) -> None:
        """
        LLM returns syntactically invalid Python, result is None.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            return_value=_stub_llm_response(
                "faster\n```python\ndef broken(\n```"
            ),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        result, cost, model_name = await generate_optimization_candidate(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
        )
        assert result is None
        assert cost is not None


@pytest.mark.asyncio
class TestOptimizePythonCode:
    """Tests for optimize_python_code (the parallel orchestrator)."""

    async def test_zero_candidates_returns_empty(self) -> None:
        """
        n_candidates=0 short-circuits to empty results.
        """
        llm_client = MagicMock()
        results, cost, code_expls, models = await optimize_python_code(
            llm_client,
            user_id="test-user",
            ctx=SingleOptimizerContext("s", "u", "x = 1\n"),
            trace_id=make_trace_id(),
            original_source_code="x = 1\n",
            n_candidates=0,
        )
        assert [] == results
        assert 0.0 == cost
        assert {} == code_expls
        assert {} == models

    async def test_multiple_candidates_parallel(self) -> None:
        """
        Multiple candidates are generated in parallel, valid ones collected.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            return_value=_stub_llm_response("faster\n```python\ny = 2\n```"),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        results, cost, code_expls, models = await optimize_python_code(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
            original_source_code="x = 1\n",
            n_candidates=3,
        )
        assert len(results) > 0
        assert cost > 0.0
        assert len(models) == len(results)

    async def test_all_failures_returns_empty(self) -> None:
        """
        All LLM calls failing returns empty results with accumulated cost.
        """
        llm_client = MagicMock()
        llm_client.call = AsyncMock(
            return_value=_stub_llm_response("no code here at all"),
        )
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        results, cost, code_expls, models = await optimize_python_code(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
            original_source_code="x = 1\n",
            n_candidates=2,
        )
        assert [] == results
        assert cost > 0.0

    async def test_model_distribution_applied(self) -> None:
        """
        With 5 candidates, both OpenAI and Anthropic models are used.
        """
        call_models: list[str] = []

        async def _capture_call(
            llm, messages, *, call_type="", trace_id="", **kw
        ):
            call_models.append(llm.name)
            return _stub_llm_response("faster\n```python\ny = 2\n```")

        llm_client = MagicMock()
        llm_client.call = AsyncMock(side_effect=_capture_call)
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        await optimize_python_code(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
            original_source_code="x = 1\n",
            n_candidates=5,
        )
        assert OPENAI_GPT_5_MINI.name in call_models
        assert ANTHROPIC_CLAUDE_SONNET_4_5.name in call_models

    async def test_line_profiler_alternation(self) -> None:
        """
        Line profiler results are passed to odd-numbered calls only.
        """
        captured_prompts: list[str] = []

        async def _capture_call(
            llm, messages, *, call_type="", trace_id="", **kw
        ):
            user_msg = next(
                m["content"] for m in messages if m["role"] == "user"
            )
            captured_prompts.append(user_msg)
            return _stub_llm_response("faster\n```python\ny = 2\n```")

        llm_client = MagicMock()
        llm_client.call = AsyncMock(side_effect=_capture_call)
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        await optimize_python_code(
            llm_client,
            user_id="test-user",
            ctx=ctx,
            trace_id=make_trace_id(),
            original_source_code="x = 1\n",
            n_candidates=4,
            line_profiler_results="Line # Hits Time ...",
        )
        has_lp = ["Line #" in p for p in captured_prompts]
        assert any(has_lp)
        assert not all(has_lp)


# ── _context.py tests ───────────────────────────────────────────────


class TestSingleOptimizerContextUserPrompt:
    """Tests for SingleOptimizerContext.get_user_prompt."""

    def test_basic_prompt_includes_source(self) -> None:
        """
        User prompt contains the source code in a markdown block.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "def foo():\n    return 1\n",
        )
        prompt = ctx.get_user_prompt("", None)
        assert "def foo():" in prompt
        assert "```python" in prompt

    def test_prompt_with_dependency_code(self) -> None:
        """
        Dependency code section is included when provided.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        prompt = ctx.get_user_prompt("import helper\n", None)
        assert "import helper" in prompt
        assert "read-only" in prompt.lower() or "READ ONLY" in prompt

    def test_prompt_with_line_profiler(self) -> None:
        """
        Line profiler section is included when results are given.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        prompt = ctx.get_user_prompt("", "Line #  Hits  Time  Per Hit")
        assert "line profiling" in prompt.lower() or "Line #" in prompt

    def test_prompt_with_runtime_info(self) -> None:
        """
        Runtime context is included when baseline_runtime_ns and loop_count given.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        prompt = ctx.get_user_prompt(
            "",
            None,
            baseline_runtime_ns=5_000_000,
            loop_count=100,
        )
        assert "5000000" in prompt
        assert "5.0ms" in prompt
        assert "100" in prompt

    def test_prompt_with_test_examples(self) -> None:
        """
        Test examples section is included when provided.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            "x = 1\n",
        )
        prompt = ctx.get_user_prompt(
            "",
            None,
            test_input_examples="assert foo(1) == 2",
        )
        assert "assert foo(1) == 2" in prompt

    def test_prompt_with_init_function(self) -> None:
        """
        Source with __init__ includes the init optimization prompt.
        """
        source = "class Foo:\n    def __init__(self):\n        self.x = 1\n"
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}{init_optimization_prompt}",
            source,
        )
        prompt = ctx.get_user_prompt("", None)
        assert "__init__" in prompt

    def test_prompt_without_init_function(self) -> None:
        """
        Source without __init__ does not include init optimization prompt.
        """
        ctx = SingleOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}{init_optimization_prompt}",
            "x = 1\n",
        )
        prompt = ctx.get_user_prompt("", None)
        assert "existing instance attributes" not in prompt


class TestMultiOptimizerContextUserPrompt:
    """Tests for MultiOptimizerContext.get_user_prompt."""

    def test_basic_prompt_includes_source(self) -> None:
        """
        Multi-file user prompt includes the original source code.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            source,
            diff_method=DiffMethod.NO_DIFF,
        )
        prompt = ctx.get_user_prompt("", None)
        assert "a.py" in prompt
        assert "b.py" in prompt

    def test_prompt_with_runtime_info(self) -> None:
        """
        Runtime context is included for multi-file context.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "Rewrite this python program to run faster.\n\n{source_code}",
            source,
            diff_method=DiffMethod.NO_DIFF,
        )
        prompt = ctx.get_user_prompt(
            "",
            None,
            baseline_runtime_ns=1_000_000_000,
            loop_count=10,
        )
        assert "1.00s" in prompt
        assert "10" in prompt


class TestMultiOptimizerContextSystemPrompt:
    """Tests for MultiOptimizerContext.get_system_prompt with diff methods."""

    def test_v4a_diff_format(self) -> None:
        """
        V4A diff method includes V4A format instructions.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "user",
            source,
            diff_method=DiffMethod.V4A,
        )
        result = ctx.get_system_prompt("3.12.9")
        assert "V4A" in result or "Begin Patch" in result

    def test_search_and_replace_format(self) -> None:
        """
        SEARCH_AND_REPLACE diff method includes S&R format instructions.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "{python_version_str} {critical_instructions}",
            "user",
            source,
            diff_method=DiffMethod.SEARCH_AND_REPLACE,
        )
        result = ctx.get_system_prompt("3.12.9")
        assert "SEARCH" in result
        assert "REPLACE" in result


class TestMultiOptimizerContextExtract:
    """Tests for MultiOptimizerContext.extract_code_and_explanation_from_llm_res."""

    def test_no_diff_extracts_files(self) -> None:
        """
        NO_DIFF method extracts markdown code blocks with file names.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "s", "u", source, diff_method=DiffMethod.NO_DIFF
        )
        content = (
            "explanation\n"
            "```python:a.py\nx = 10\n```\n"
            "```python:b.py\ny = 20\n```"
        )
        result = ctx.extract_code_and_explanation_from_llm_res(content)
        assert "explanation" == result.explanation
        assert "a.py" in result.code
        assert "b.py" in result.code


class TestMultiOptimizerContextValidate:
    """Tests for MultiOptimizerContext.validate_and_parse_source_code."""

    def test_valid_multi_file(self) -> None:
        """
        Valid multi-file source passes validation.
        """
        source = "```python:a.py\nx = 1\n```\n```python:b.py\ny = 2\n```"
        ctx = MultiOptimizerContext(
            "s", "u", source, diff_method=DiffMethod.NO_DIFF
        )
        ctx.validate_and_parse_source_code(source, feature_version=(3, 12))

    def test_invalid_multi_file_syntax_error(self) -> None:
        """
        Invalid Python in any file raises SyntaxError.
        """
        source = "```python:a.py\ndef broken(\n```"
        ctx = MultiOptimizerContext(
            "s", "u", source, diff_method=DiffMethod.NO_DIFF
        )
        with pytest.raises(SyntaxError):
            ctx.validate_and_parse_source_code(source, feature_version=(3, 12))


class TestMultiOptimizerContextIsValidCode:
    """Tests for MultiOptimizerContext.is_valid_code edge cases."""

    def test_valid_code_with_matching_files(self) -> None:
        """
        Extracted files matching originals pass validation.
        """
        source = "```python:a.py\nx = 1\n```"
        ctx = MultiOptimizerContext(
            "s", "u", source, diff_method=DiffMethod.NO_DIFF
        )
        ctx.extracted_code_and_expl = CodeStrAndExplanation(
            code="```python:a.py\nx = 2\n```",
            explanation="faster",
        )
        assert ctx.is_valid_code()

    def test_none_extracted_is_invalid(self) -> None:
        """
        No extracted code is invalid.
        """
        source = "```python:a.py\nx = 1\n```"
        ctx = MultiOptimizerContext(
            "s", "u", source, diff_method=DiffMethod.NO_DIFF
        )
        assert not ctx.is_valid_code()


# ── _router.py tests ────────────────────────────────────────────────


class TestValidateRequest:
    """Tests for _validate_request helper."""

    def test_empty_source_code_raises_400(self) -> None:
        """
        Empty source_code field raises HTTPException with 400.
        """
        from fastapi import HTTPException

        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="",
            trace_id=make_trace_id(),
            python_version="3.12.9",
        )
        ctx = BaseOptimizerContext("s", "u", "")
        with pytest.raises(HTTPException) as exc_info:
            _validate_request(req, ctx)
        assert 400 == exc_info.value.status_code

    def test_invalid_trace_id_raises_400(self) -> None:
        """
        Invalid trace_id raises HTTPException with 400.
        """
        from fastapi import HTTPException

        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="x = 1",
            trace_id="not-a-uuid",
            python_version="3.12.9",
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(HTTPException) as exc_info:
            _validate_request(req, ctx)
        assert 400 == exc_info.value.status_code

    def test_missing_python_version_raises_400(self) -> None:
        """
        None python_version raises HTTPException with 400.
        """
        from fastapi import HTTPException

        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="x = 1",
            trace_id=make_trace_id(),
            python_version=None,
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(HTTPException) as exc_info:
            _validate_request(req, ctx)
        assert 400 == exc_info.value.status_code

    def test_invalid_python_version_raises_400(self) -> None:
        """
        Bad python_version format raises HTTPException with 400.
        """
        from fastapi import HTTPException

        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="x = 1",
            trace_id=make_trace_id(),
            python_version="2.7.18",
        )
        ctx = BaseOptimizerContext("s", "u", "x = 1")
        with pytest.raises(HTTPException) as exc_info:
            _validate_request(req, ctx)
        assert 400 == exc_info.value.status_code

    def test_invalid_source_code_syntax_raises_400(self) -> None:
        """
        Source code with syntax errors raises HTTPException with 400.
        """
        from fastapi import HTTPException

        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="def broken(",
            trace_id=make_trace_id(),
            python_version="3.12.9",
        )
        ctx = SingleOptimizerContext("s", "u", "def broken(")
        with pytest.raises(HTTPException) as exc_info:
            _validate_request(req, ctx)
        assert 400 == exc_info.value.status_code

    def test_valid_request_returns_version_tuple(self) -> None:
        """
        Valid request returns parsed python version tuple.
        """
        from codeflash_api.optimize._router import _validate_request

        req = OptimizeRequest(
            source_code="x = 1\n",
            trace_id=make_trace_id(),
            python_version="3.12.9",
        )
        ctx = SingleOptimizerContext("s", "u", "x = 1\n")
        result = _validate_request(req, ctx)
        assert (3, 12, 9) == result


@pytest.mark.asyncio
class TestOptimizeEndpoint:
    """Tests for POST /ai/optimize endpoint."""

    async def test_success_returns_optimizations(
        self, client, mock_llm_client
    ) -> None:
        """
        Valid request with working LLM returns optimization results.
        """
        mock_llm_client.call = AsyncMock(
            return_value=_stub_llm_response("faster\n```python\ny = 2\n```"),
        )
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "x = 1\n",
                "trace_id": make_trace_id(),
                "python_version": "3.12.9",
                "n_candidates": 1,
            },
        )
        assert 200 == resp.status_code
        body = resp.json()
        assert "optimizations" in body
        assert len(body["optimizations"]) > 0
        assert "source_code" in body["optimizations"][0]

    async def test_empty_source_returns_400(self, client) -> None:
        """
        Empty source_code returns 400 error.
        """
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "",
                "trace_id": make_trace_id(),
                "python_version": "3.12.9",
            },
        )
        assert 400 == resp.status_code

    async def test_invalid_trace_id_returns_400(self, client) -> None:
        """
        Invalid trace_id returns 400 error.
        """
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "x = 1\n",
                "trace_id": "bad-uuid",
                "python_version": "3.12.9",
            },
        )
        assert 400 == resp.status_code

    async def test_no_optimizations_returns_422(
        self, client, mock_llm_client
    ) -> None:
        """
        LLM returning no valid code produces 422.
        """
        mock_llm_client.call = AsyncMock(
            return_value=_stub_llm_response("no code here"),
        )
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "x = 1\n",
                "trace_id": make_trace_id(),
                "python_version": "3.12.9",
                "n_candidates": 1,
            },
        )
        assert 422 == resp.status_code
        body = resp.json()
        assert "detail" in body

    async def test_async_flag_uses_async_prompts(
        self, client, mock_llm_client
    ) -> None:
        """
        is_async=True uses the async system prompt.
        """
        captured_messages: list[list[dict]] = []

        async def _capture_call(llm, messages, **kw):
            captured_messages.append(messages)
            return _stub_llm_response("faster\n```python\ny = 2\n```")

        mock_llm_client.call = AsyncMock(side_effect=_capture_call)
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "x = 1\n",
                "trace_id": make_trace_id(),
                "python_version": "3.12.9",
                "n_candidates": 1,
                "is_async": True,
            },
        )
        assert 200 == resp.status_code
        assert len(captured_messages) > 0

    async def test_missing_python_version_returns_400(self, client) -> None:
        """
        Missing python_version returns 400.
        """
        resp = await client.post(
            "/ai/optimize",
            json={
                "source_code": "x = 1\n",
                "trace_id": make_trace_id(),
            },
        )
        assert 400 == resp.status_code
