from __future__ import annotations

import json
import uuid
from unittest.mock import AsyncMock

import pytest

from codeflash_api.languages.python._markdown import (
    wrap_code_in_markdown,
)
from codeflash_api.llm._client import (
    LLMOutputUnparseableError,
    LLMResponse,
    LLMUsage,
)
from codeflash_api.review._router import (
    _build_review_messages,
    _parse_review_response,
    get_optimization_review,
)
from codeflash_api.review.schemas import (
    OptimizationReviewErrorResponse,
    OptimizationReviewRequest,
    OptimizationReviewResponse,
    ReviewLevel,
)


def _make_request(**overrides: object) -> OptimizationReviewRequest:
    """
    Build a minimal valid OptimizationReviewRequest with overrides.
    """
    defaults: dict[str, object] = {
        "trace_id": str(uuid.uuid4()),
        "original_code": "def f(): pass",
        "optimized_code": "def f(): return None",
        "original_runtime": "100ms",
        "optimized_runtime": "10ms",
        "speedup": "90%",
        "coverage_message": "100%",
        "loop_count": 1000,
        "explanation": "Faster return.",
        "calling_fn_details": "main() calls f()",
    }
    defaults.update(overrides)
    return OptimizationReviewRequest(**defaults)  # type: ignore[arg-type]


def _make_llm_json(
    rating: str = "high",
    *,
    performance_score: int = 8,
    code_quality_score: int = 7,
    hot_path_score: int = 6,
) -> str:
    """
    Build a realistic LLM response with a JSON code block.
    """
    json_block = json.dumps(
        {
            "rating": rating,
            "performance_score": performance_score,
            "performance_reasoning": "Good speedup",
            "code_quality_score": code_quality_score,
            "code_quality_reasoning": "Clean code",
            "hot_path_score": hot_path_score,
            "hot_path_reasoning": "Called in a loop",
        }
    )
    return (
        f"Analysis of the optimization.\n\n"
        f"```json\n{json_block}\n```\n\n"
        f"Overall looks good."
    )


class TestReviewLevel:
    """Tests for ReviewLevel enum."""

    def test_values(self) -> None:
        """
        All three levels are accessible.
        """
        assert "high" == ReviewLevel.HIGH.value
        assert "medium" == ReviewLevel.MEDIUM.value
        assert "low" == ReviewLevel.LOW.value

    def test_from_string(self) -> None:
        """
        Enum can be constructed from a string value.
        """
        assert ReviewLevel.HIGH == ReviewLevel("high")


class TestOptimizationReviewSchemas:
    """Tests for Pydantic review schemas."""

    def test_request_defaults(self) -> None:
        """
        Default fields are set correctly.
        """
        req = _make_request()
        assert "python" == req.language
        assert "0.18.2" == req.codeflash_version
        assert req.code_diff is None

    def test_request_all_fields(self) -> None:
        """
        All optional fields can be set.
        """
        req = _make_request(
            existing_tests="test_1",
            generated_tests="test_2",
            replay_tests="test_3",
            benchmark_details="details",
            python_version="3.12.0",
            call_sequence=1,
            acceptance_reason="runtime",
            original_concurrency_ratio=1.0,
            optimized_concurrency_ratio=10.0,
            original_throughput=100,
            optimized_throughput=1000,
        )
        assert "test_1" == req.existing_tests
        assert 10.0 == req.optimized_concurrency_ratio

    def test_response_model(self) -> None:
        """
        Response is serializable with all fields.
        """
        resp = OptimizationReviewResponse(
            review=ReviewLevel.HIGH,
            review_explanation="Looks good",
            performance_score=8,
            code_quality_score=7,
            hot_path_score=6,
        )
        data = resp.model_dump()
        assert "high" == data["review"]
        assert 8 == data["performance_score"]

    def test_response_optional_scores(self) -> None:
        """
        Scores default to None.
        """
        resp = OptimizationReviewResponse(
            review=ReviewLevel.LOW,
            review_explanation="Bad",
        )
        assert resp.performance_score is None
        assert resp.code_quality_score is None
        assert resp.hot_path_score is None

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


class TestBuildReviewMessages:
    """Tests for _build_review_messages."""

    def test_with_original_and_optimized(self) -> None:
        """
        Messages are built from original + optimized code.
        """
        data = _make_request()
        messages = _build_review_messages(data)
        assert 2 == len(messages)
        assert "system" == messages[0]["role"]
        assert "user" == messages[1]["role"]
        assert "def f(): pass" in messages[1]["content"]

    def test_with_code_diff_legacy(self) -> None:
        """
        Legacy code_diff path builds inline XML prompt.
        """
        data = _make_request(
            original_code=None,
            optimized_code=None,
            code_diff="--- a/f.py\n+++ b/f.py\n@@ ...",
        )
        messages = _build_review_messages(data)
        assert 2 == len(messages)
        assert "code_changes" in messages[1]["content"]
        assert "--- a/f.py" in messages[1]["content"]

    def test_missing_both_raises(self) -> None:
        """
        Missing both code paths raises ValueError.
        """
        data = _make_request(
            original_code=None,
            optimized_code=None,
            code_diff=None,
        )
        with pytest.raises(ValueError, match="Must provide"):
            _build_review_messages(data)

    def test_old_version_wraps_generated_tests(self) -> None:
        """
        Versions <= 0.18.2 wrap generated_tests in markdown.
        """
        data = _make_request(
            codeflash_version="0.18.2",
            generated_tests="assert True",
        )
        _build_review_messages(data)
        assert data.generated_tests == wrap_code_in_markdown("assert True")

    def test_new_version_skips_wrapping(self) -> None:
        """
        Versions > 0.18.2 leave generated_tests unchanged.
        """
        data = _make_request(
            codeflash_version="0.19.0",
            generated_tests="assert True",
        )
        _build_review_messages(data)
        assert "assert True" == data.generated_tests

    def test_language_version_python(self) -> None:
        """
        Python version is included in the prompt.
        """
        data = _make_request(python_version="3.12.0")
        messages = _build_review_messages(data)
        assert "3.12.0" in messages[1]["content"]

    def test_language_version_non_python(self) -> None:
        """
        Non-Python uses 'Not specified' for version.
        """
        data = _make_request(language="javascript")
        messages = _build_review_messages(data)
        assert "Not specified" in messages[1]["content"]

    def test_optional_fields_default_to_not_available(
        self,
    ) -> None:
        """
        None optional fields render as 'Not Available'.
        """
        data = _make_request(
            existing_tests=None,
            generated_tests=None,
        )
        messages = _build_review_messages(data)
        assert "Not Available" in messages[1]["content"]

    def test_acceptance_reason_defaults_to_runtime(self) -> None:
        """
        Missing acceptance_reason defaults to 'runtime'.
        """
        data = _make_request(acceptance_reason=None)
        messages = _build_review_messages(data)
        content = messages[1]["content"]
        assert "runtime" in content


class TestParseReviewResponse:
    """Tests for _parse_review_response."""

    def test_valid_high(self) -> None:
        """
        Valid JSON block with high rating is parsed.
        """
        text = _make_llm_json("high")
        resp = _parse_review_response(text, cost=0.1)
        assert ReviewLevel.HIGH == resp.review
        assert 8 == resp.performance_score
        assert 7 == resp.code_quality_score
        assert 6 == resp.hot_path_score
        assert "Analysis" in resp.review_explanation
        assert "Overall looks good" in resp.review_explanation

    def test_valid_low(self) -> None:
        """
        Low rating is parsed correctly.
        """
        text = _make_llm_json("low")
        resp = _parse_review_response(text, cost=0.0)
        assert ReviewLevel.LOW == resp.review

    def test_scores_appended_to_explanation(self) -> None:
        """
        Dimension scores are appended to the explanation.
        """
        text = _make_llm_json()
        resp = _parse_review_response(text, cost=0.0)
        assert "Dimension Scores" in resp.review_explanation
        assert "Performance: 8/10" in resp.review_explanation
        assert "Code Quality: 7/10" in resp.review_explanation
        assert "Hot Path: 6/10" in resp.review_explanation

    def test_no_json_block_raises(self) -> None:
        """
        Missing JSON block raises LLMOutputUnparseableError.
        """
        with pytest.raises(LLMOutputUnparseableError):
            _parse_review_response("no json here", cost=0.5)

    def test_invalid_json_raises(self) -> None:
        """
        Malformed JSON raises LLMOutputUnparseableError.
        """
        text = "before\n```json\n{bad json}\n```\nafter"
        with pytest.raises(LLMOutputUnparseableError):
            _parse_review_response(text, cost=0.5)

    def test_missing_rating_key_raises(self) -> None:
        """
        JSON without 'rating' raises LLMOutputUnparseableError.
        """
        text = '```json\n{"performance_score": 8}\n```'
        with pytest.raises(LLMOutputUnparseableError):
            _parse_review_response(text, cost=0.0)

    def test_invalid_rating_value_raises(self) -> None:
        """
        Unknown rating value raises LLMOutputUnparseableError.
        """
        text = '```json\n{"rating": "excellent"}\n```'
        with pytest.raises(LLMOutputUnparseableError):
            _parse_review_response(text, cost=0.0)

    def test_no_scores_no_summary(self) -> None:
        """
        When all scores are absent, no summary is appended.
        """
        json_block = json.dumps({"rating": "medium"})
        text = f"Good.\n```json\n{json_block}\n```\nEnd."
        resp = _parse_review_response(text, cost=0.0)
        assert "Dimension Scores" not in resp.review_explanation
        assert resp.performance_score is None

    def test_cost_preserved_on_error(self) -> None:
        """
        Error includes the accumulated cost.
        """
        with pytest.raises(LLMOutputUnparseableError) as exc_info:
            _parse_review_response("garbage", cost=1.23)
        assert 1.23 == exc_info.value.cost


@pytest.mark.asyncio
class TestGetOptimizationReview:
    """Tests for get_optimization_review."""

    async def test_success_first_attempt(self) -> None:
        """
        Successful LLM call on first attempt returns the review.
        """
        llm_client = AsyncMock()
        llm_client.call.return_value = LLMResponse(
            content=_make_llm_json("high"),
            usage=LLMUsage(input_tokens=100, output_tokens=50),
            raw_response=None,
            cost=0.1,
        )
        data = _make_request()

        resp, cost = await get_optimization_review(llm_client, data)
        assert ReviewLevel.HIGH == resp.review
        assert 0.1 == cost
        assert 1 == llm_client.call.call_count

    async def test_retry_on_unparseable(self) -> None:
        """
        Unparseable first attempt retries and succeeds.
        """
        llm_client = AsyncMock()
        llm_client.call.side_effect = [
            LLMResponse(
                content="no json here",
                usage=LLMUsage(input_tokens=100, output_tokens=50),
                raw_response=None,
                cost=0.05,
            ),
            LLMResponse(
                content=_make_llm_json("medium"),
                usage=LLMUsage(input_tokens=100, output_tokens=50),
                raw_response=None,
                cost=0.05,
            ),
        ]
        data = _make_request()

        resp, cost = await get_optimization_review(llm_client, data)
        assert ReviewLevel.MEDIUM == resp.review
        assert 2 == llm_client.call.call_count
        assert cost == pytest.approx(0.10)

    async def test_all_retries_exhausted(self) -> None:
        """
        All retries unparseable raises LLMOutputUnparseableError.
        """
        llm_client = AsyncMock()
        llm_client.call.return_value = LLMResponse(
            content="garbage",
            usage=LLMUsage(input_tokens=100, output_tokens=50),
            raw_response=None,
            cost=0.05,
        )
        data = _make_request()

        with pytest.raises(LLMOutputUnparseableError):
            await get_optimization_review(llm_client, data)
        assert 2 == llm_client.call.call_count

    async def test_llm_client_exception_propagates(
        self,
    ) -> None:
        """
        Non-parse errors from the LLM client propagate.
        """
        llm_client = AsyncMock()
        llm_client.call.side_effect = RuntimeError("connection")
        data = _make_request()

        with pytest.raises(RuntimeError, match="connection"):
            await get_optimization_review(llm_client, data)

    async def test_value_error_propagates(self) -> None:
        """
        ValueError from missing code paths propagates.
        """
        llm_client = AsyncMock()
        data = _make_request(
            original_code=None,
            optimized_code=None,
            code_diff=None,
        )

        with pytest.raises(ValueError, match="Must provide"):
            await get_optimization_review(llm_client, data)
