from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import pytest

from codeflash_api.llm._client import (
    LLMClient,
    LLMOutputUnparseableError,
    LLMResponse,
    LLMUsage,
)
from codeflash_api.llm._cost import calculate_llm_cost
from codeflash_api.llm._models import (
    ANTHROPIC_CLAUDE_HAIKU_4_5,
    ANTHROPIC_CLAUDE_SONNET_4_5,
    LLM,
    OPENAI_GPT_4_1,
    OPENAI_GPT_5_MINI,
)
from codeflash_api.llm._retry import (
    ANTHROPIC_MAX_INPUT_TOKENS,
    CHARS_PER_TOKEN_ESTIMATE,
    TRANSIENT_LLM_ERRORS,
)


class TestLLMModel:
    """Tests for LLM model definitions."""

    def test_frozen(self) -> None:
        """
        LLM instances are immutable.
        """
        model = LLM(name="test", model_type="openai")
        with pytest.raises(AttributeError):
            model.name = "changed"

    def test_cost_defaults(self) -> None:
        """
        Cost rates default to None.
        """
        model = LLM(name="test", model_type="openai")
        assert model.input_cost is None
        assert model.cached_input_cost is None
        assert model.output_cost is None

    def test_openai_gpt_4_1_rates(self) -> None:
        """
        GPT-4.1 has correct cost rates.
        """
        assert 2.00 == OPENAI_GPT_4_1.input_cost
        assert 0.50 == OPENAI_GPT_4_1.cached_input_cost
        assert 8.00 == OPENAI_GPT_4_1.output_cost

    def test_openai_gpt_5_mini_rates(self) -> None:
        """
        GPT-5-mini has correct cost rates.
        """
        assert 0.25 == OPENAI_GPT_5_MINI.input_cost
        assert 0.03 == OPENAI_GPT_5_MINI.cached_input_cost
        assert 2.00 == OPENAI_GPT_5_MINI.output_cost

    def test_anthropic_sonnet_rates(self) -> None:
        """
        Claude Sonnet 4.5 has correct cost rates.
        """
        assert 3.00 == ANTHROPIC_CLAUDE_SONNET_4_5.input_cost
        assert ANTHROPIC_CLAUDE_SONNET_4_5.cached_input_cost is None
        assert 15.00 == ANTHROPIC_CLAUDE_SONNET_4_5.output_cost

    def test_anthropic_haiku_rates(self) -> None:
        """
        Claude Haiku 4.5 has correct cost rates.
        """
        assert 1.00 == ANTHROPIC_CLAUDE_HAIKU_4_5.input_cost
        assert 5.00 == ANTHROPIC_CLAUDE_HAIKU_4_5.output_cost

    def test_all_models_have_names(self) -> None:
        """
        Every predefined model has a non-empty name.
        """
        for model in (
            OPENAI_GPT_4_1,
            OPENAI_GPT_5_MINI,
            ANTHROPIC_CLAUDE_SONNET_4_5,
            ANTHROPIC_CLAUDE_HAIKU_4_5,
        ):
            assert model.name


class TestCostCalculation:
    """Tests for calculate_llm_cost."""

    def test_no_usage_returns_zero(self) -> None:
        """
        Response without usage attribute returns 0.0.
        """
        assert 0.0 == calculate_llm_cost(object(), OPENAI_GPT_4_1)

    def test_openai_basic(self) -> None:
        """
        OpenAI cost calculated from prompt/completion tokens.
        """
        usage = SimpleNamespace(
            prompt_tokens=1000,
            completion_tokens=500,
            prompt_tokens_details=None,
        )
        response = SimpleNamespace(usage=usage)
        cost = calculate_llm_cost(response, OPENAI_GPT_4_1)
        expected = (1000 * 2.00 + 500 * 8.00) / 1_000_000
        assert pytest.approx(expected) == cost

    def test_openai_with_cache(self) -> None:
        """
        OpenAI cached tokens use cached_input_cost rate.
        """
        details = SimpleNamespace(cached_tokens=400)
        usage = SimpleNamespace(
            prompt_tokens=1000,
            completion_tokens=200,
            prompt_tokens_details=details,
        )
        response = SimpleNamespace(usage=usage)
        cost = calculate_llm_cost(response, OPENAI_GPT_4_1)
        expected = (600 * 2.00 + 400 * 0.50 + 200 * 8.00) / 1_000_000
        assert pytest.approx(expected) == cost

    def test_anthropic_basic(self) -> None:
        """
        Anthropic cost uses input_tokens + output_tokens.
        """
        usage = SimpleNamespace(
            input_tokens=1000,
            output_tokens=500,
            cache_read_input_tokens=0,
            cache_creation_input_tokens=0,
        )
        response = SimpleNamespace(usage=usage)
        cost = calculate_llm_cost(response, ANTHROPIC_CLAUDE_SONNET_4_5)
        expected = (1000 * 3.00 + 500 * 15.00) / 1_000_000
        assert pytest.approx(expected) == cost

    def test_anthropic_with_cache_creation(self) -> None:
        """
        Cache creation tokens count toward input cost.
        """
        usage = SimpleNamespace(
            input_tokens=800,
            output_tokens=200,
            cache_read_input_tokens=100,
            cache_creation_input_tokens=300,
        )
        response = SimpleNamespace(usage=usage)
        cost = calculate_llm_cost(response, ANTHROPIC_CLAUDE_SONNET_4_5)
        non_cached = 800 + 300
        expected = (non_cached * 3.00 + 100 * 3.00 + 200 * 15.00) / 1_000_000
        assert pytest.approx(expected) == cost

    def test_model_without_costs(self) -> None:
        """
        Model with None costs treats them as 0.
        """
        model = LLM(name="free", model_type="openai")
        usage = SimpleNamespace(
            prompt_tokens=1000,
            completion_tokens=500,
            prompt_tokens_details=None,
        )
        response = SimpleNamespace(usage=usage)
        assert 0.0 == calculate_llm_cost(response, model)


class TestRetryConstants:
    """Tests for retry module constants."""

    def test_transient_errors_count(self) -> None:
        """
        Eight transient error types are defined.
        """
        assert 8 == len(TRANSIENT_LLM_ERRORS)

    def test_all_are_exception_subclasses(self) -> None:
        """
        Every transient error is an Exception subclass.
        """
        for exc_type in TRANSIENT_LLM_ERRORS:
            assert issubclass(exc_type, Exception)

    def test_constants(self) -> None:
        """
        Token limit and char estimate are set.
        """
        assert 195_000 == ANTHROPIC_MAX_INPUT_TOKENS
        assert 4 == CHARS_PER_TOKEN_ESTIMATE


class TestLLMUsage:
    """Tests for LLMUsage."""

    def test_frozen(self) -> None:
        """
        LLMUsage is immutable.
        """
        usage = LLMUsage(input_tokens=10, output_tokens=20)
        with pytest.raises(AttributeError):
            usage.input_tokens = 99


class TestLLMResponse:
    """Tests for LLMResponse."""

    def test_cost_defaults_to_zero(self) -> None:
        """
        Cost defaults to 0.0.
        """
        resp = LLMResponse(
            content="hello",
            usage=LLMUsage(input_tokens=1, output_tokens=1),
            raw_response=None,
        )
        assert 0.0 == resp.cost

    def test_frozen(self) -> None:
        """
        LLMResponse is immutable.
        """
        resp = LLMResponse(
            content="hello",
            usage=LLMUsage(input_tokens=1, output_tokens=1),
            raw_response=None,
        )
        with pytest.raises(AttributeError):
            resp.content = "changed"


class TestLLMOutputUnparseableError:
    """Tests for LLMOutputUnparseableError."""

    def test_default_message(self) -> None:
        """
        Default message is set.
        """
        err = LLMOutputUnparseableError()
        assert "could not be parsed" in str(err)
        assert 0.0 == err.cost

    def test_custom_cost(self) -> None:
        """
        Cost is preserved on the exception.
        """
        err = LLMOutputUnparseableError(cost=1.23)
        assert 1.23 == err.cost


class TestLLMClient:
    """Tests for LLMClient."""

    async def test_unsupported_model_type(self) -> None:
        """
        Unknown model_type raises ValueError.
        """
        client = LLMClient()
        bad_model = LLM.__new__(LLM)
        object.__setattr__(bad_model, "name", "bad")
        object.__setattr__(bad_model, "model_type", "unknown")
        object.__setattr__(bad_model, "input_cost", None)
        object.__setattr__(bad_model, "cached_input_cost", None)
        object.__setattr__(bad_model, "output_cost", None)

        with (
            patch.object(LLMClient, "_ensure_clients"),
            pytest.raises(ValueError, match="Unsupported"),
        ):
            await client.call(bad_model, [])

    async def test_anthropic_not_configured(self) -> None:
        """
        Calling Anthropic without configuration raises RuntimeError.
        """
        client = LLMClient()
        with pytest.raises(RuntimeError, match="not configured"):
            await client._call_anthropic(ANTHROPIC_CLAUDE_SONNET_4_5, [], 1024)

    async def test_openai_not_configured(self) -> None:
        """
        Calling OpenAI without configuration raises RuntimeError.
        """
        client = LLMClient()
        with pytest.raises(RuntimeError, match="not configured"):
            await client._call_openai(OPENAI_GPT_4_1, [], 1024)

    async def test_anthropic_prompt_too_large(self) -> None:
        """
        Oversized prompt raises ValueError.
        """
        client = LLMClient()
        client._anthropic_client = AsyncMock()
        huge = "x" * (
            ANTHROPIC_MAX_INPUT_TOKENS * CHARS_PER_TOKEN_ESTIMATE + 100
        )
        messages = [{"role": "user", "content": huge}]

        with pytest.raises(ValueError, match="too large"):
            await client._call_anthropic(
                ANTHROPIC_CLAUDE_SONNET_4_5, messages, 1024
            )

    async def test_anthropic_empty_response(self) -> None:
        """
        Empty Anthropic response raises LLMOutputUnparseableError.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        response = SimpleNamespace(
            content=[],
            usage=SimpleNamespace(
                input_tokens=10,
                output_tokens=0,
                cache_read_input_tokens=0,
                cache_creation_input_tokens=0,
            ),
        )
        mock_client.messages.create = AsyncMock(return_value=response)
        client._anthropic_client = mock_client

        with pytest.raises(LLMOutputUnparseableError):
            await client._call_anthropic(
                ANTHROPIC_CLAUDE_SONNET_4_5,
                [{"role": "user", "content": "hi"}],
                1024,
            )

    async def test_anthropic_success(self) -> None:
        """
        Successful Anthropic call returns LLMResponse.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        block = SimpleNamespace(text="hello world")
        response = SimpleNamespace(
            content=[block],
            usage=SimpleNamespace(
                input_tokens=10,
                output_tokens=5,
                cache_read_input_tokens=0,
                cache_creation_input_tokens=0,
            ),
        )
        mock_client.messages.create = AsyncMock(return_value=response)
        client._anthropic_client = mock_client

        result = await client._call_anthropic(
            ANTHROPIC_CLAUDE_SONNET_4_5,
            [{"role": "user", "content": "hi"}],
            1024,
        )
        assert "hello world" == result.content
        assert 10 == result.usage.input_tokens
        assert 5 == result.usage.output_tokens

    async def test_anthropic_separates_system_prompt(self) -> None:
        """
        System messages are extracted into the system kwarg.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        block = SimpleNamespace(text="ok")
        response = SimpleNamespace(
            content=[block],
            usage=SimpleNamespace(
                input_tokens=10,
                output_tokens=5,
                cache_read_input_tokens=0,
                cache_creation_input_tokens=0,
            ),
        )
        mock_client.messages.create = AsyncMock(return_value=response)
        client._anthropic_client = mock_client

        await client._call_anthropic(
            ANTHROPIC_CLAUDE_SONNET_4_5,
            [
                {"role": "system", "content": "be helpful"},
                {"role": "user", "content": "hi"},
            ],
            1024,
        )

        call_kwargs = mock_client.messages.create.call_args[1]
        assert "be helpful" == call_kwargs["system"]
        assert all(m["role"] != "system" for m in call_kwargs["messages"])

    async def test_openai_empty_response(self) -> None:
        """
        Empty OpenAI response raises LLMOutputUnparseableError.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        response = SimpleNamespace(
            choices=[],
            usage=SimpleNamespace(
                prompt_tokens=10,
                completion_tokens=0,
                prompt_tokens_details=None,
            ),
        )
        mock_client.chat.completions.create = AsyncMock(return_value=response)
        client._openai_client = mock_client

        with pytest.raises(LLMOutputUnparseableError):
            await client._call_openai(
                OPENAI_GPT_4_1,
                [{"role": "user", "content": "hi"}],
                1024,
            )

    async def test_openai_success(self) -> None:
        """
        Successful OpenAI call returns LLMResponse.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        response = SimpleNamespace(
            choices=[
                SimpleNamespace(message=SimpleNamespace(content="hello"))
            ],
            usage=SimpleNamespace(
                prompt_tokens=10,
                completion_tokens=5,
                prompt_tokens_details=None,
            ),
        )
        mock_client.chat.completions.create = AsyncMock(return_value=response)
        client._openai_client = mock_client

        result = await client._call_openai(
            OPENAI_GPT_4_1,
            [{"role": "user", "content": "hi"}],
            1024,
        )
        assert "hello" == result.content
        assert 10 == result.usage.input_tokens
        assert 5 == result.usage.output_tokens

    async def test_openai_gpt5_mini_uses_max_completion_tokens(
        self,
    ) -> None:
        """
        GPT-5-mini uses max_completion_tokens instead of max_tokens.
        """
        client = LLMClient()
        mock_client = AsyncMock()
        response = SimpleNamespace(
            choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))],
            usage=SimpleNamespace(
                prompt_tokens=10,
                completion_tokens=5,
                prompt_tokens_details=None,
            ),
        )
        mock_client.chat.completions.create = AsyncMock(return_value=response)
        client._openai_client = mock_client

        await client._call_openai(
            OPENAI_GPT_5_MINI,
            [{"role": "user", "content": "hi"}],
            2048,
        )

        call_kwargs = mock_client.chat.completions.create.call_args[1]
        assert 2048 == call_kwargs["max_completion_tokens"]
        assert "max_tokens" not in call_kwargs
