from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from codeflash_api.ranking._router import (
    SCORING_DIMENSIONS,
    _compute_weighted_score,
    _extract_json_from_response,
    _parse_json_response,
    _parse_scores_text,
    _scores_to_ranking,
    rank_optimizations,
)
from codeflash_api.ranking.schemas import (
    RankErrorResponse,
    RankRequest,
    RankResponse,
)


class TestRankRequest:
    """Tests for the RankRequest schema."""

    def test_minimal_fields(self) -> None:
        """
        Only trace_id, speedups, diffs, and optimization_ids are required.
        """
        req = RankRequest(
            trace_id="abc",
            speedups=[1.5],
            diffs=["diff"],
            optimization_ids=["opt1"],
        )
        assert "abc" == req.trace_id
        assert req.python_version is None
        assert req.function_references is None

    def test_all_fields(self) -> None:
        """
        All optional fields are accepted.
        """
        req = RankRequest(
            trace_id="abc",
            speedups=[1.5, 2.0],
            diffs=["d1", "d2"],
            optimization_ids=["o1", "o2"],
            python_version="3.12.0",
            function_references="some refs",
        )
        assert "3.12.0" == req.python_version
        assert "some refs" == req.function_references


class TestRankResponse:
    """Tests for the RankResponse schema."""

    def test_without_scores(self) -> None:
        """
        Scores field is optional and defaults to None.
        """
        resp = RankResponse(
            ranking=[0, 1],
            explanation="test",
        )
        assert resp.scores is None

    def test_with_scores(self) -> None:
        """
        Scores can include per-candidate dimension scores.
        """
        resp = RankResponse(
            ranking=[0, 1],
            explanation="test",
            scores={
                1: {"speedup_quality": 8.0, "readability": 7.0},
            },
        )
        assert 1 in resp.scores


class TestRankErrorResponse:
    """Tests for the RankErrorResponse schema."""

    def test_error_message(self) -> None:
        """
        Error response contains the error string.
        """
        resp = RankErrorResponse(error="something failed")
        assert "something failed" == resp.error


class TestExtractJsonFromResponse:
    """Tests for _extract_json_from_response."""

    def test_direct_json(self) -> None:
        """
        Raw JSON string is parsed directly.
        """
        raw = json.dumps({"ranking": [1, 2], "explanation": "ok"})
        result = _extract_json_from_response(raw)
        assert result is not None
        assert [1, 2] == result["ranking"]

    def test_json_in_markdown_code_block(self) -> None:
        """
        JSON inside a ```json code block is extracted.
        """
        text = (
            "Here is the result:\n"
            "```json\n"
            '{"ranking": [2, 1], "explanation": "done"}\n'
            "```\n"
        )
        result = _extract_json_from_response(text)
        assert result is not None
        assert [2, 1] == result["ranking"]

    def test_json_with_surrounding_text(self) -> None:
        """
        JSON embedded in plain text is found by brace matching.
        """
        text = 'Some preamble. {"ranking": [1], "explanation": "x"} trailing.'
        result = _extract_json_from_response(text)
        assert result is not None
        assert [1] == result["ranking"]

    def test_no_json(self) -> None:
        """
        Returns None when no JSON can be extracted.
        """
        assert _extract_json_from_response("no json here") is None

    def test_empty_string(self) -> None:
        """
        Returns None for empty input.
        """
        assert _extract_json_from_response("") is None

    def test_invalid_json_in_code_block(self) -> None:
        """
        Returns None when code block contains invalid JSON.
        """
        text = "```json\n{invalid json}\n```"
        assert _extract_json_from_response(text) is None

    def test_nested_json_braces(self) -> None:
        """
        Properly matched nested braces are handled.
        """
        text = (
            'text {"scores": {"1": {"a": 1}}, '
            '"ranking": [1], "explanation": "ok"} more'
        )
        result = _extract_json_from_response(text)
        assert result is not None
        assert [1] == result["ranking"]


class TestParseJsonResponse:
    """Tests for _parse_json_response."""

    def test_valid_two_candidates(self) -> None:
        """
        A valid response with scores, ranking, and explanation
        for two candidates is parsed.
        """
        content = json.dumps(
            {
                "scores": {
                    "1": {
                        "speedup_quality": 7,
                        "readability": 8,
                        "diff_precision": 9,
                        "correctness_confidence": 8,
                    },
                    "2": {
                        "speedup_quality": 9,
                        "readability": 6,
                        "diff_precision": 7,
                        "correctness_confidence": 7,
                    },
                },
                "ranking": [2, 1],
                "explanation": "Candidate 2 is better",
            }
        )
        result = _parse_json_response(content, 2)
        assert result is not None
        ranking, scores, explanation = result
        assert [2, 1] == ranking
        assert scores is not None
        assert 2 == len(scores)
        assert "Candidate 2 is better" == explanation

    def test_invalid_ranking_values(self) -> None:
        """
        Ranking with wrong candidate numbers returns None.
        """
        content = json.dumps(
            {
                "ranking": [3, 1],
                "explanation": "bad",
            }
        )
        assert _parse_json_response(content, 2) is None

    def test_missing_ranking_key(self) -> None:
        """
        Response without 'ranking' key returns None.
        """
        content = json.dumps({"explanation": "no ranking"})
        assert _parse_json_response(content, 1) is None

    def test_scores_clamped_to_range(self) -> None:
        """
        Scores outside 1-10 are clamped to the valid range.
        """
        content = json.dumps(
            {
                "scores": {
                    "1": {
                        "speedup_quality": 15,
                        "readability": -2,
                        "diff_precision": 5,
                        "correctness_confidence": 7,
                    },
                },
                "ranking": [1],
                "explanation": "clamped",
            }
        )
        result = _parse_json_response(content, 1)
        assert result is not None
        _, scores, _ = result
        assert scores is not None
        assert 10.0 == scores[1]["speedup_quality"]
        assert 1.0 == scores[1]["readability"]

    def test_missing_dimension_in_scores(self) -> None:
        """
        Candidate missing a dimension gets scores set to None.
        """
        content = json.dumps(
            {
                "scores": {
                    "1": {
                        "speedup_quality": 7,
                        "readability": 8,
                    },
                },
                "ranking": [1],
                "explanation": "incomplete",
            }
        )
        result = _parse_json_response(content, 1)
        assert result is not None
        _, scores, _ = result
        assert scores is None

    def test_not_json(self) -> None:
        """
        Non-JSON input returns None.
        """
        assert _parse_json_response("not json", 2) is None

    def test_ranking_not_a_list(self) -> None:
        """
        Ranking as a non-list type returns None.
        """
        content = json.dumps(
            {
                "ranking": "1,2",
                "explanation": "bad",
            }
        )
        assert _parse_json_response(content, 2) is None


class TestParseScoresText:
    """Tests for _parse_scores_text."""

    def test_valid_two_candidates(self) -> None:
        """
        Valid text format with all dimensions parses correctly.
        """
        text = (
            "1: speedup_quality=7, readability=8,"
            " diff_precision=9, correctness_confidence=8\n"
            "2: speedup_quality=5, readability=6,"
            " diff_precision=4, correctness_confidence=3\n"
        )
        result = _parse_scores_text(text, 2)
        assert result is not None
        assert 7.0 == result[1]["speedup_quality"]
        assert 3.0 == result[2]["correctness_confidence"]

    def test_missing_candidate(self) -> None:
        """
        Returns None when not all candidates have scores.
        """
        text = (
            "1: speedup_quality=7, readability=8,"
            " diff_precision=9, correctness_confidence=8\n"
        )
        assert _parse_scores_text(text, 2) is None

    def test_missing_dimension(self) -> None:
        """
        Returns None when a candidate is missing a dimension.
        """
        text = "1: speedup_quality=7, readability=8\n"
        assert _parse_scores_text(text, 1) is None

    def test_empty_input(self) -> None:
        """
        Returns None for empty input.
        """
        assert _parse_scores_text("", 1) is None

    def test_clamped_values(self) -> None:
        """
        Scores are clamped to the 1-10 range.
        """
        text = (
            "1: speedup_quality=0, readability=11,"
            " diff_precision=5, correctness_confidence=7\n"
        )
        result = _parse_scores_text(text, 1)
        assert result is not None
        assert 1.0 == result[1]["speedup_quality"]
        assert 10.0 == result[1]["readability"]


class TestComputeWeightedScore:
    """Tests for _compute_weighted_score."""

    def test_all_tens(self) -> None:
        """
        All dimensions at 10 produces a weighted sum of 10.
        """
        scores = dict.fromkeys(SCORING_DIMENSIONS, 10.0)
        assert 10.0 == _compute_weighted_score(scores)

    def test_all_ones(self) -> None:
        """
        All dimensions at 1 produces a weighted sum of 1.
        """
        scores = dict.fromkeys(SCORING_DIMENSIONS, 1.0)
        assert 1.0 == _compute_weighted_score(scores)

    def test_missing_dimension_defaults_to_five(self) -> None:
        """
        A missing dimension defaults to 5.0 in the calculation.
        """
        scores = {
            "speedup_quality": 10.0,
            "readability": 10.0,
            "diff_precision": 10.0,
        }
        expected = 10.0 * 0.30 + 10.0 * 0.25 + 10.0 * 0.20 + 5.0 * 0.25
        assert expected == pytest.approx(_compute_weighted_score(scores))


class TestScoresToRanking:
    """Tests for _scores_to_ranking."""

    def test_higher_score_ranks_first(self) -> None:
        """
        The candidate with the higher weighted score ranks first.
        """
        scores = {
            1: dict.fromkeys(SCORING_DIMENSIONS, 5.0),
            2: dict.fromkeys(SCORING_DIMENSIONS, 9.0),
        }
        assert [2, 1] == _scores_to_ranking(scores)

    def test_single_candidate(self) -> None:
        """
        A single candidate produces a single-element ranking.
        """
        scores = {
            1: dict.fromkeys(SCORING_DIMENSIONS, 7.0),
        }
        assert [1] == _scores_to_ranking(scores)

    def test_three_candidates(self) -> None:
        """
        Three candidates are ordered by weighted score descending.
        """
        scores = {
            1: dict.fromkeys(SCORING_DIMENSIONS, 3.0),
            2: dict.fromkeys(SCORING_DIMENSIONS, 8.0),
            3: dict.fromkeys(SCORING_DIMENSIONS, 5.0),
        }
        assert [2, 3, 1] == _scores_to_ranking(scores)


class TestRankOptimizations:
    """Tests for rank_optimizations."""

    @pytest.mark.asyncio
    async def test_json_response_parsed(self) -> None:
        """
        A well-formed JSON response produces a 1-indexed ranking.
        """
        llm_response = MagicMock()
        llm_response.content = json.dumps(
            {
                "scores": {
                    "1": {
                        "speedup_quality": 5,
                        "readability": 5,
                        "diff_precision": 5,
                        "correctness_confidence": 5,
                    },
                    "2": {
                        "speedup_quality": 9,
                        "readability": 9,
                        "diff_precision": 9,
                        "correctness_confidence": 9,
                    },
                },
                "ranking": [2, 1],
                "explanation": "Candidate 2 is superior",
            }
        )
        llm_response.cost = 0.01

        mock_client = AsyncMock()
        mock_client.call.return_value = llm_response

        data = RankRequest(
            trace_id="test-trace",
            speedups=[1.5, 3.0],
            diffs=["diff1", "diff2"],
            optimization_ids=["opt1", "opt2"],
        )

        ranking, scores, explanation, cost = await rank_optimizations(
            mock_client, "user1", data
        )

        assert [2, 1] == ranking
        assert scores is not None
        assert "Candidate 2 is superior" == explanation
        assert 0.01 == cost
        mock_client.call.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_legacy_xml_fallback(self) -> None:
        """
        Legacy XML-tag format is parsed when JSON fails.
        """
        llm_response = MagicMock()
        llm_response.content = (
            "<rank>2,1</rank><explain>Two is better</explain>"
        )
        llm_response.cost = 0.02

        mock_client = AsyncMock()
        mock_client.call.return_value = llm_response

        data = RankRequest(
            trace_id="test-trace",
            speedups=[1.0, 2.0],
            diffs=["d1", "d2"],
            optimization_ids=["o1", "o2"],
        )

        ranking, scores, explanation, cost = await rank_optimizations(
            mock_client, "user1", data
        )

        assert [2, 1] == ranking
        assert scores is None
        assert "Two is better" == explanation

    @pytest.mark.asyncio
    async def test_unparseable_raises_422(self) -> None:
        """
        Completely unparseable LLM output raises HTTPException 422.
        """
        from fastapi import HTTPException

        llm_response = MagicMock()
        llm_response.content = "totally unparseable garbage"
        llm_response.cost = 0.0

        mock_client = AsyncMock()
        mock_client.call.return_value = llm_response

        data = RankRequest(
            trace_id="test-trace",
            speedups=[1.0, 2.0],
            diffs=["d1", "d2"],
            optimization_ids=["o1", "o2"],
        )

        with pytest.raises(HTTPException) as exc_info:
            await rank_optimizations(mock_client, "user1", data)

        assert 422 == exc_info.value.status_code

    @pytest.mark.asyncio
    async def test_scores_derive_ranking(self) -> None:
        """
        When XML scores are present but ranking is invalid,
        ranking is derived from scores.
        """
        scores_block = (
            "1: speedup_quality=3, readability=3,"
            " diff_precision=3, correctness_confidence=3\n"
            "2: speedup_quality=9, readability=9,"
            " diff_precision=9, correctness_confidence=9\n"
        )
        llm_response = MagicMock()
        llm_response.content = (
            f"<rank>99,99</rank>"
            f"<scores>{scores_block}</scores>"
            f"<explain>derived</explain>"
        )
        llm_response.cost = 0.0

        mock_client = AsyncMock()
        mock_client.call.return_value = llm_response

        data = RankRequest(
            trace_id="test-trace",
            speedups=[1.0, 2.0],
            diffs=["d1", "d2"],
            optimization_ids=["o1", "o2"],
        )

        ranking, scores, explanation, _cost = await rank_optimizations(
            mock_client, "user1", data
        )

        assert [2, 1] == ranking
        assert scores is not None
        assert "derived" == explanation
