"""Tests for _refinement — refinement, repair, and adaptive."""

from __future__ import annotations

from typing import Any

import attrs
import pytest

from codeflash_core.exceptions import (
    AIServiceConnectionError,
    AIServiceError,
)
from codeflash_python.ai._refinement import (
    AdaptiveCandidate,
    AdaptiveOptimizeRequest,
    CodeRepairRequest,
    OptimizedCandidateSource,
    RefinementRequest,
    _parse_candidate,
    adaptive_optimize,
    code_repair,
    optimize_code_refinement,
)


class MockClient:
    """Mock AIClient with a configurable post() method."""

    def __init__(
        self,
        post_return: dict[str, Any] | None = None,
        post_side_effect: Exception | None = None,
    ) -> None:
        self._post_return = post_return or {}
        self._post_side_effect = post_side_effect
        self.last_endpoint = ""
        self.last_payload: Any = None

    def post(
        self,
        endpoint: str,
        payload: dict[str, Any] | list[Any],
    ) -> dict[str, Any]:
        """Record the call and return the configured response."""
        self.last_endpoint = endpoint
        self.last_payload = payload
        if self._post_side_effect:
            raise self._post_side_effect
        return self._post_return


def make_refinement_request(
    **overrides: Any,
) -> RefinementRequest:
    """Build a RefinementRequest with sensible defaults."""
    defaults: dict[str, Any] = {
        "optimization_id": "opt-1",
        "original_source_code": "def foo(): pass",
        "read_only_dependency_code": "",
        "original_code_runtime": 1000,
        "optimized_source_code": "def foo(): return 1",
        "optimized_explanation": "Optimized",
        "optimized_code_runtime": 500,
        "speedup": "2x",
        "trace_id": "trace-1",
        "original_line_profiler_results": "",
        "optimized_line_profiler_results": "",
    }
    defaults.update(overrides)
    return RefinementRequest(**defaults)


def make_code_repair_request(
    **overrides: Any,
) -> CodeRepairRequest:
    """Build a CodeRepairRequest with sensible defaults."""
    defaults: dict[str, Any] = {
        "optimization_id": "opt-1",
        "original_source_code": "def foo(): pass",
        "modified_source_code": "def foo(): return 1",
        "trace_id": "trace-1",
        "test_diffs": (),
    }
    defaults.update(overrides)
    return CodeRepairRequest(**defaults)


def make_adaptive_request(
    **overrides: Any,
) -> AdaptiveOptimizeRequest:
    """Build an AdaptiveOptimizeRequest with sensible defaults."""
    candidate = AdaptiveCandidate(
        optimization_id="opt-1",
        source_code="def foo(): return 1",
        explanation="Faster",
        source=OptimizedCandidateSource.OPTIMIZE,
        speedup="2x",
    )
    defaults: dict[str, Any] = {
        "trace_id": "trace-1",
        "original_source_code": "def foo(): pass",
        "candidates": (candidate,),
    }
    defaults.update(overrides)
    return AdaptiveOptimizeRequest(**defaults)


class TestOptimizedCandidateSource:
    """OptimizedCandidateSource str enum."""

    def test_all_values_present(self) -> None:
        """All six enum members exist."""
        members = {m.name for m in OptimizedCandidateSource}
        expected = {
            "OPTIMIZE",
            "OPTIMIZE_LP",
            "REFINE",
            "REPAIR",
            "ADAPTIVE",
            "JIT_REWRITE",
        }
        assert expected == members

    def test_values_are_strings(self) -> None:
        """Each member value is a string."""
        for member in OptimizedCandidateSource:
            assert isinstance(member.value, str)

    def test_is_str_subclass(self) -> None:
        """Members are instances of str."""
        for member in OptimizedCandidateSource:
            assert isinstance(member, str)


class TestRefinementRequest:
    """RefinementRequest frozen attrs class."""

    def test_construction(self) -> None:
        """Can construct with all required fields."""
        req = make_refinement_request()

        assert "opt-1" == req.optimization_id
        assert "def foo(): pass" == req.original_source_code
        assert "" == req.read_only_dependency_code
        assert 1000 == req.original_code_runtime
        assert "def foo(): return 1" == req.optimized_source_code
        assert "Optimized" == req.optimized_explanation
        assert 500 == req.optimized_code_runtime
        assert "2x" == req.speedup
        assert "trace-1" == req.trace_id

    def test_frozen(self) -> None:
        """Raises on attribute assignment."""
        req = make_refinement_request()

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            req.optimization_id = "changed"  # type: ignore[misc]

    def test_optional_field_defaults_to_none(self) -> None:
        """function_references defaults to None."""
        req = make_refinement_request()

        assert req.function_references is None


class TestCodeRepairRequest:
    """CodeRepairRequest frozen attrs class."""

    def test_construction(self) -> None:
        """Can construct with all required fields."""
        req = make_code_repair_request()

        assert "opt-1" == req.optimization_id
        assert "def foo(): pass" == req.original_source_code
        assert "def foo(): return 1" == req.modified_source_code
        assert "trace-1" == req.trace_id
        assert () == req.test_diffs

    def test_frozen(self) -> None:
        """Raises on attribute assignment."""
        req = make_code_repair_request()

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            req.trace_id = "changed"  # type: ignore[misc]


class TestAdaptiveCandidate:
    """AdaptiveCandidate frozen attrs class."""

    def test_construction(self) -> None:
        """Can construct with all fields."""
        candidate = AdaptiveCandidate(
            optimization_id="opt-1",
            source_code="def foo(): return 1",
            explanation="Faster",
            source=OptimizedCandidateSource.OPTIMIZE,
            speedup="2x",
        )

        assert "opt-1" == candidate.optimization_id
        assert "def foo(): return 1" == candidate.source_code
        assert "Faster" == candidate.explanation
        assert OptimizedCandidateSource.OPTIMIZE == candidate.source
        assert "2x" == candidate.speedup

    def test_frozen(self) -> None:
        """Raises on attribute assignment."""
        candidate = AdaptiveCandidate(
            optimization_id="opt-1",
            source_code="def foo(): return 1",
            explanation="Faster",
            source=OptimizedCandidateSource.OPTIMIZE,
            speedup="2x",
        )

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            candidate.source_code = "changed"  # type: ignore[misc]


class TestAdaptiveOptimizeRequest:
    """AdaptiveOptimizeRequest frozen attrs class."""

    def test_construction(self) -> None:
        """Can construct with all required fields."""
        candidate = AdaptiveCandidate(
            optimization_id="opt-1",
            source_code="def foo(): return 1",
            explanation="Faster",
            source=OptimizedCandidateSource.OPTIMIZE,
            speedup="2x",
        )
        req = AdaptiveOptimizeRequest(
            trace_id="trace-1",
            original_source_code="def foo(): pass",
            candidates=(candidate,),
        )

        assert "trace-1" == req.trace_id
        assert "def foo(): pass" == req.original_source_code
        assert 1 == len(req.candidates)
        assert candidate is req.candidates[0]

    def test_frozen(self) -> None:
        """Raises on attribute assignment."""
        req = make_adaptive_request()

        with pytest.raises(attrs.exceptions.FrozenInstanceError):
            req.trace_id = "changed"  # type: ignore[misc]


class TestOptimizeCodeRefinement:
    """optimize_code_refinement AI service call."""

    def test_successful_response(self) -> None:
        """Successful response returns list of Candidate objects."""
        client = MockClient(
            post_return={
                "refinements": [
                    {
                        "source_code": "def foo(): return 2",
                        "explanation": "Refined",
                        "optimization_id": "opt-1",
                    },
                ],
            },
        )

        result = optimize_code_refinement(
            client,  # type: ignore[arg-type]
            [make_refinement_request()],
        )

        assert "/refinement" == client.last_endpoint
        assert 1 == len(result)
        assert "def foo(): return 2" == result[0].code
        assert "Refined" == result[0].explanation
        assert "opt-1" == result[0].candidate_id

    def test_http_error_raises_ai_service_error(self) -> None:
        """HTTP error raises AIServiceError."""
        client = MockClient(
            post_side_effect=AIServiceError(500, "fail"),
        )

        with pytest.raises(AIServiceError):
            optimize_code_refinement(
                client,  # type: ignore[arg-type]
                [make_refinement_request()],
            )

    def test_connection_error_raises_connection_error(self) -> None:
        """Connection error raises AIServiceConnectionError."""
        client = MockClient(
            post_side_effect=AIServiceConnectionError("refused"),
        )

        with pytest.raises(AIServiceConnectionError):
            optimize_code_refinement(
                client,  # type: ignore[arg-type]
                [make_refinement_request()],
            )

    def test_empty_requests_returns_empty(self) -> None:
        """Empty refinements list returns empty list."""
        client = MockClient(post_return={"refinements": []})

        result = optimize_code_refinement(
            client,  # type: ignore[arg-type]
            [],
        )

        assert [] == result

    def test_empty_source_code_filtered_out(self) -> None:
        """Response with empty source_code is filtered out."""
        client = MockClient(
            post_return={
                "refinements": [
                    {
                        "source_code": "",
                        "explanation": "Empty",
                        "optimization_id": "opt-1",
                    },
                    {
                        "source_code": "def foo(): return 3",
                        "explanation": "Valid",
                        "optimization_id": "opt-2",
                    },
                ],
            },
        )

        result = optimize_code_refinement(
            client,  # type: ignore[arg-type]
            [make_refinement_request()],
        )

        assert 1 == len(result)
        assert "def foo(): return 3" == result[0].code


class TestCodeRepair:
    """code_repair AI service call."""

    def test_successful_response(self) -> None:
        """Successful response returns a Candidate."""
        client = MockClient(
            post_return={
                "source_code": "def foo(): return fixed",
                "explanation": "Fixed",
                "optimization_id": "opt-1",
            },
        )

        result = code_repair(
            client,  # type: ignore[arg-type]
            make_code_repair_request(),
        )

        assert "/code_repair" == client.last_endpoint
        assert result is not None
        assert "def foo(): return fixed" == result.code
        assert "Fixed" == result.explanation
        assert "opt-1" == result.candidate_id

    def test_empty_source_code_returns_none(self) -> None:
        """Returns None when response has empty source_code."""
        client = MockClient(
            post_return={
                "source_code": "",
                "explanation": "Empty",
                "optimization_id": "opt-1",
            },
        )

        result = code_repair(
            client,  # type: ignore[arg-type]
            make_code_repair_request(),
        )

        assert result is None

    def test_http_error_raises_ai_service_error(self) -> None:
        """HTTP error raises AIServiceError."""
        client = MockClient(
            post_side_effect=AIServiceError(500, "fail"),
        )

        with pytest.raises(AIServiceError):
            code_repair(
                client,  # type: ignore[arg-type]
                make_code_repair_request(),
            )

    def test_connection_error_raises_connection_error(self) -> None:
        """Connection error raises AIServiceConnectionError."""
        client = MockClient(
            post_side_effect=AIServiceConnectionError("refused"),
        )

        with pytest.raises(AIServiceConnectionError):
            code_repair(
                client,  # type: ignore[arg-type]
                make_code_repair_request(),
            )


class TestAdaptiveOptimize:
    """adaptive_optimize AI service call."""

    def test_successful_response(self) -> None:
        """Successful response returns a Candidate."""
        client = MockClient(
            post_return={
                "source_code": "def foo(): return adaptive",
                "explanation": "Adaptive",
                "optimization_id": "opt-1",
            },
        )

        result = adaptive_optimize(
            client,  # type: ignore[arg-type]
            make_adaptive_request(),
        )

        assert "/adaptive_optimize" == client.last_endpoint
        assert result is not None
        assert "def foo(): return adaptive" == result.code
        assert "Adaptive" == result.explanation
        assert "opt-1" == result.candidate_id

    def test_empty_source_code_returns_none(self) -> None:
        """Returns None when response has empty source_code."""
        client = MockClient(
            post_return={
                "source_code": "",
                "explanation": "Empty",
                "optimization_id": "opt-1",
            },
        )

        result = adaptive_optimize(
            client,  # type: ignore[arg-type]
            make_adaptive_request(),
        )

        assert result is None

    def test_http_error_raises_ai_service_error(self) -> None:
        """HTTP error raises AIServiceError."""
        client = MockClient(
            post_side_effect=AIServiceError(500, "fail"),
        )

        with pytest.raises(AIServiceError):
            adaptive_optimize(
                client,  # type: ignore[arg-type]
                make_adaptive_request(),
            )

    def test_connection_error_raises_connection_error(self) -> None:
        """Connection error raises AIServiceConnectionError."""
        client = MockClient(
            post_side_effect=AIServiceConnectionError("refused"),
        )

        with pytest.raises(AIServiceConnectionError):
            adaptive_optimize(
                client,  # type: ignore[arg-type]
                make_adaptive_request(),
            )


class TestParseCandidateMarkdown:
    """_parse_candidate markdown code-block parsing."""

    def test_plain_code_passes_through(self) -> None:
        """Plain code without markdown fences is used as-is."""
        data = {
            "source_code": "def foo(): return 1",
            "explanation": "Plain",
            "optimization_id": "opt-1",
        }

        result = _parse_candidate(data, "REPAIR")

        assert result is not None
        assert "def foo(): return 1" == result.code

    def test_markdown_fenced_code_is_extracted(self) -> None:
        """Markdown-fenced code block is parsed and extracted."""
        data = {
            "source_code": "```python\ndef foo(): return 2\n```",
            "explanation": "Fenced",
            "optimization_id": "opt-2",
        }

        result = _parse_candidate(data, "REFINE")

        assert result is not None
        assert "def foo(): return 2" == result.code

    def test_markdown_with_file_path_is_extracted(self) -> None:
        """Markdown block with file path suffix is parsed."""
        data = {
            "source_code": ("```python:src/mod.py\ndef bar(): pass\n```"),
            "explanation": "With path",
            "optimization_id": "opt-3",
        }

        result = _parse_candidate(data, "REFINE")

        assert result is not None
        assert "def bar(): pass" == result.code

    def test_multiple_code_blocks_joined(self) -> None:
        """Multiple markdown blocks are joined with double newlines."""
        data = {
            "source_code": (
                "```python\ndef a(): pass\n```\n```python\ndef b(): pass\n```"
            ),
            "explanation": "Multi",
            "optimization_id": "opt-4",
        }

        result = _parse_candidate(data, "ADAPTIVE")

        assert result is not None
        assert "def a(): pass" in result.code
        assert "def b(): pass" in result.code
        assert "\n\n" in result.code

    def test_empty_source_code_returns_none(self) -> None:
        """Empty source_code returns None."""
        data = {
            "source_code": "",
            "explanation": "Empty",
            "optimization_id": "opt-5",
        }

        assert _parse_candidate(data, "REPAIR") is None

    def test_missing_source_code_returns_none(self) -> None:
        """Missing source_code key returns None."""
        data: dict[str, Any] = {
            "explanation": "No code",
            "optimization_id": "opt-6",
        }

        assert _parse_candidate(data, "REPAIR") is None

    def test_source_field_is_forwarded(self) -> None:
        """The source parameter is set on the Candidate."""
        data = {
            "source_code": "x = 1",
            "explanation": "Test",
            "optimization_id": "opt-7",
        }

        result = _parse_candidate(data, "ADAPTIVE")

        assert result is not None
        assert "ADAPTIVE" == result.source
