"""Tests for the LLM gateway package."""

from __future__ import annotations

import asyncio
import json
from typing import Any

import httpx
import pytest

from llm_gateway.config import GatewayConfig, ProviderConfig, TaskClass
from llm_gateway.exceptions import BudgetExceeded, SchemaValidationError
from llm_gateway.gateway import Gateway, TelemetryRow, run_task
from llm_gateway.mock import get_mock_output, MOCK_OUTPUTS


# ---- Fixtures ----


def mock_config(**overrides: Any) -> GatewayConfig:
    """Build a config in mock mode (no API key)."""
    primary = ProviderConfig(
        name="primary",
        base_url="https://mock.example.com/v1",
        api_key="",
        model="glm-5.2",
    )
    defaults: dict[str, Any] = {
        "primary": primary,
        "fallback": None,
        "cheap_model": "glm-5.2",
        "strong_model": "glm-5.2",
        "budgets": {"score": 2000, "extract": 4000, "default": 4000},
        "max_retries": 2,
    }
    defaults.update(overrides)
    return GatewayConfig(**defaults)


def real_config(**overrides: Any) -> GatewayConfig:
    """Build a config with a fake API key (non-mock mode)."""
    primary = ProviderConfig(
        name="primary",
        base_url="https://mock.example.com/v1",
        api_key="fake-key-1234",
        model="glm-5.2",
    )
    defaults: dict[str, Any] = {
        "primary": primary,
        "fallback": None,
        "cheap_model": "glm-5.2",
        "strong_model": "glm-5.2",
        "budgets": {"score": 2000, "extract": 4000, "default": 4000},
        "max_retries": 2,
    }
    defaults.update(overrides)
    return GatewayConfig(**defaults)


# ---- Mock mode tests ----


class TestMockMode:
    async def test_mock_returns_deterministic_output(self) -> None:
        """Mock mode returns deterministic canned outputs per task."""
        config = mock_config()
        gw = Gateway(config)
        result_a = await gw.run_task("score", "Score this job")
        result_b = await gw.run_task("score", "Score this job")
        assert result_a == result_b
        assert result_a["score"] == 75
        await gw.aclose()

    async def test_mock_different_tasks_different_output(self) -> None:
        config = mock_config()
        gw = Gateway(config)
        score_result = await gw.run_task("score", "prompt")
        extract_result = await gw.run_task("extract", "prompt")
        assert score_result != extract_result
        assert "score" in score_result
        assert "company" in extract_result
        await gw.aclose()

    async def test_mock_unknown_task(self) -> None:
        config = mock_config()
        gw = Gateway(config)
        result = await gw.run_task("unknown_task", "prompt")
        assert result["result"] == "mock output"
        assert result["task"] == "unknown_task"
        await gw.aclose()

    async def test_mock_mode_property(self) -> None:
        config = mock_config()
        assert config.mock_mode is True

    async def test_real_mode_not_mock(self) -> None:
        config = real_config()
        assert config.mock_mode is False


# ---- Schema validation tests ----


class TestSchemaValidation:
    async def test_schema_passes(self) -> None:
        config = mock_config()
        gw = Gateway(config)
        schema = {
            "type": "object",
            "properties": {
                "score": {"type": "number"},
                "rationale": {"type": "object"},
            },
            "required": ["score"],
        }
        result = await gw.run_task("score", "prompt", schema=schema)
        assert "score" in result
        await gw.aclose()

    async def test_schema_fails(self) -> None:
        """Schema validation failure should raise SchemaValidationError."""
        config = mock_config()
        gw = Gateway(config)
        # The mock output for 'score' has score=75 (number). We require a string,
        # which should fail validation.
        bad_schema = {
            "type": "object",
            "properties": {
                "score": {"type": "string"},
            },
            "required": ["score"],
        }
        with pytest.raises(SchemaValidationError):
            await gw.run_task("score", "prompt", schema=bad_schema)
        await gw.aclose()

    async def test_schema_missing_required_field(self) -> None:
        config = mock_config()
        gw = Gateway(config)
        schema = {
            "type": "object",
            "required": ["nonexistent_field"],
        }
        with pytest.raises(SchemaValidationError):
            await gw.run_task("score", "prompt", schema=schema)
        await gw.aclose()


# ---- Budget guard tests ----


class TestBudgetGuard:
    async def test_budget_exceeded_raises_before_call(self) -> None:
        """Over-budget prompt should raise BudgetExceeded before any call."""
        config = mock_config(budgets={"score": 10})
        gw = Gateway(config)
        # 10 token budget, ~4 chars/token, so >40 chars should exceed.
        long_prompt = "x" * 100
        with pytest.raises(BudgetExceeded):
            await gw.run_task("score", long_prompt)
        await gw.aclose()

    async def test_budget_within_limit_does_not_raise(self) -> None:
        config = mock_config(budgets={"score": 10000})
        gw = Gateway(config)
        result = await gw.run_task("score", "short prompt")
        assert result["score"] == 75
        await gw.aclose()

    async def test_budget_guard_in_real_mode(self) -> None:
        """Budget guard must raise before call even in real (non-mock) mode."""
        config = real_config(budgets={"score": 10})
        gw = Gateway(config)
        with pytest.raises(BudgetExceeded):
            await gw.run_task("score", "x" * 100)
        await gw.aclose()

    async def test_default_budget_fallback(self) -> None:
        """Unknown task should use default budget."""
        config = mock_config(budgets={"score": 2000, "default": 100})
        gw = Gateway(config)
        # Unknown task uses default=100, so >400 chars exceeds.
        with pytest.raises(BudgetExceeded):
            await gw.run_task("unknown_task", "x" * 500)
        await gw.aclose()


# ---- Telemetry tests ----


class TestTelemetry:
    async def test_telemetry_sink_called_in_mock_mode(self) -> None:
        sink_calls: list[TelemetryRow] = []

        async def sink(row: TelemetryRow) -> None:
            sink_calls.append(row)

        config = mock_config()
        gw = Gateway(config, telemetry_sink=sink)
        await gw.run_task("score", "test prompt")
        assert len(sink_calls) == 1
        assert sink_calls[0].task == "score"
        assert sink_calls[0].mock is True
        await gw.aclose()

    async def test_telemetry_sink_sync_callable(self) -> None:
        """Sync sinks should also work (no await needed)."""
        sink_calls: list[TelemetryRow] = []

        def sync_sink(row: TelemetryRow) -> None:
            sink_calls.append(row)

        config = mock_config()
        gw = Gateway(config, telemetry_sink=sync_sink)
        await gw.run_task("score", "test prompt")
        assert len(sink_calls) == 1
        await gw.aclose()

    async def test_no_sink_no_error(self) -> None:
        config = mock_config()
        gw = Gateway(config, telemetry_sink=None)
        result = await gw.run_task("score", "prompt")
        assert result["score"] == 75
        await gw.aclose()

    async def test_telemetry_row_to_dict(self) -> None:
        row = TelemetryRow(task="score", model="glm-5.2", provider="mock")
        d = row.to_dict()
        assert d["task"] == "score"
        assert d["model"] == "glm-5.2"
        assert d["provider"] == "mock"
        assert "id" in d


# ---- Config tests ----


class TestGatewayConfig:
    def test_from_env_mock_mode(self) -> None:
        env = {}
        config = GatewayConfig.from_env(env=env)
        assert config.mock_mode is True
        assert config.primary.model == "glm-5.2"

    def test_from_env_real_mode(self) -> None:
        env = {
            "LLM_PRIMARY_KEY": "test-key",
            "LLM_PRIMARY_BASE_URL": "https://api.example.com/v1",
            "LLM_PRIMARY_MODEL": "custom-model",
        }
        config = GatewayConfig.from_env(env=env)
        assert config.mock_mode is False
        assert config.primary.api_key == "test-key"
        assert config.primary.model == "custom-model"

    def test_from_env_ollama_key(self) -> None:
        env = {"OLLAMA_API_KEY": "ollama-key-123"}
        config = GatewayConfig.from_env(env=env)
        assert config.mock_mode is False
        assert config.primary.api_key == "ollama-key-123"

    def test_from_env_budgets(self) -> None:
        env = {"LLM_BUDGET_SCORE": "500"}
        config = GatewayConfig.from_env(env=env)
        assert config.budgets["score"] == 500

    def test_task_class_routing(self) -> None:
        config = mock_config()
        assert config.get_task_class("score") == TaskClass.CHEAP
        assert config.get_task_class("extract") == TaskClass.CHEAP
        assert config.get_task_class("cv_assist") == TaskClass.CHEAP
        assert config.get_task_class("email_classify") == TaskClass.CHEAP
        assert config.get_task_class("deadline_extract") == TaskClass.CHEAP
        assert config.get_task_class("critique") == TaskClass.STRONG
        assert config.get_task_class("cl_critique") == TaskClass.STRONG
        assert config.get_task_class("cv_tailor") == TaskClass.STRONG

    def test_get_model_routing(self) -> None:
        config = mock_config(cheap_model="cheap-model", strong_model="strong-model")
        assert config.get_model("score") == "cheap-model"
        assert config.get_model("critique") == "strong-model"

    def test_paid_fallback_detection(self) -> None:
        """Config should detect paid provider URLs."""
        env = {
            "LLM_PRIMARY_KEY": "key",
            "LLM_FALLBACK_BASE_URL": "https://api.openai.com/v1",
            "LLM_FALLBACK_KEY": "fb-key",
        }
        config = GatewayConfig.from_env(env=env)
        assert config.fallback is not None
        assert config.fallback.is_paid is True

    def test_paid_fallback_not_used_for_cheap(self) -> None:
        """The gateway must not route cheap tasks to paid fallback.

        We verify this by checking _select_provider returns primary for cheap.
        """
        primary = ProviderConfig(
            name="primary", base_url="https://a.com/v1", api_key="k", model="m"
        )
        fallback = ProviderConfig(
            name="fallback",
            base_url="https://api.openai.com/v1",
            api_key="k2",
            model="m2",
            is_paid=True,
        )
        config = GatewayConfig(primary=primary, fallback=fallback)
        gw = Gateway(config)
        provider = gw._select_provider("score")
        assert provider.name == "primary"
        # The _select_provider method enforces this by always returning primary
        # for cheap tasks, never the paid fallback.

    def test_non_paid_fallback(self) -> None:
        """Non-paid (e.g. ollama) fallback is fine."""
        env = {
            "LLM_PRIMARY_KEY": "key",
            "LLM_FALLBACK_BASE_URL": "https://api.ollama-cloud.com/v1",
            "LLM_FALLBACK_KEY": "fb-key",
        }
        config = GatewayConfig.from_env(env=env)
        assert config.fallback is not None
        assert config.fallback.is_paid is False


# ---- Provider call tests (with mocked HTTP) ----


class TestProviderCalls:
    async def test_real_mode_calls_provider(self) -> None:
        """In non-mock mode, the gateway should make an HTTP call."""
        config = real_config()
        mock_response_data = {
            "choices": [
                {
                    "message": {
                        "content": json.dumps({"score": 85, "rationale": {"ok": True}})
                    }
                }
            ],
            "usage": {"prompt_tokens": 50, "completion_tokens": 30},
            "model": "glm-5.2",
        }

        def handler(request: httpx.Request) -> httpx.Response:
            return httpx.Response(200, json=mock_response_data)

        client = httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://mock.example.com",
        )
        gw = Gateway(config, http_client=client)
        result = await gw.run_task("score", "Score this job")
        assert result["score"] == 85
        assert "_usage" not in result  # _usage should be popped
        await client.aclose()

    async def test_retry_on_429(self) -> None:
        """Gateway should retry on 429 then succeed."""
        config = real_config(max_retries=2)
        call_count = 0

        def handler(request: httpx.Request) -> httpx.Response:
            nonlocal call_count
            call_count += 1
            if call_count < 2:
                return httpx.Response(429, json={"error": "rate limited"})
            return httpx.Response(
                200,
                json={
                    "choices": [
                        {"message": {"content": json.dumps({"score": 50})}}
                    ],
                    "usage": {"prompt_tokens": 10, "completion_tokens": 5},
                },
            )

        client = httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://mock.example.com",
        )
        gw = Gateway(config, http_client=client)
        result = await gw.run_task("score", "test")
        assert result["score"] == 50
        assert call_count == 2
        await client.aclose()

    async def test_retry_exhausted_raises(self) -> None:
        """After all retries, ProviderError should be raised."""
        config = real_config(max_retries=1)

        def handler(request: httpx.Request) -> httpx.Response:
            return httpx.Response(500, json={"error": "server error"})

        client = httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
            base_url="https://mock.example.com",
        )
        from llm_gateway.exceptions import ProviderError

        gw = Gateway(config, http_client=client)
        with pytest.raises(ProviderError):
            await gw.run_task("critique", "test")
        await client.aclose()

    async def test_fallback_used_for_strong_task(self) -> None:
        """Strong tasks should fall back when primary fails."""
        primary = ProviderConfig(
            name="primary",
            base_url="https://primary.example.com/v1",
            api_key="pk",
            model="glm-5.2",
        )
        fallback = ProviderConfig(
            name="fallback",
            base_url="https://fallback.example.com/v1",
            api_key="fk",
            model="glm-5.2",
            is_paid=False,
        )
        config = GatewayConfig(
            primary=primary,
            fallback=fallback,
            budgets={"critique": 4000, "default": 4000},
            max_retries=0,
        )

        def handler(request: httpx.Request) -> httpx.Response:
            if "primary.example.com" in str(request.url):
                return httpx.Response(500, json={"error": "primary down"})
            return httpx.Response(
                200,
                json={
                    "choices": [
                        {"message": {"content": json.dumps({"comments": []})}}
                    ],
                    "usage": {"prompt_tokens": 20, "completion_tokens": 10},
                    "model": "glm-5.2",
                },
            )

        client = httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
        )
        gw = Gateway(config, http_client=client)
        result = await gw.run_task("critique", "review this")
        assert result == {"comments": []}
        await client.aclose()

    async def test_no_fallback_for_cheap_task(self) -> None:
        """Cheap tasks must not use fallback even when primary fails."""
        primary = ProviderConfig(
            name="primary",
            base_url="https://primary.example.com/v1",
            api_key="pk",
            model="glm-5.2",
        )
        fallback = ProviderConfig(
            name="fallback",
            base_url="https://fallback.example.com/v1",
            api_key="fk",
            model="glm-5.2",
        )
        config = GatewayConfig(
            primary=primary,
            fallback=fallback,
            budgets={"score": 4000, "default": 4000},
            max_retries=0,
        )

        def handler(request: httpx.Request) -> httpx.Response:
            return httpx.Response(500, json={"error": "down"})

        client = httpx.AsyncClient(
            transport=httpx.MockTransport(handler),
        )
        from llm_gateway.exceptions import ProviderError

        gw = Gateway(config, http_client=client)
        with pytest.raises(ProviderError):
            await gw.run_task("score", "score this")
        await client.aclose()


# ---- Convenience function test ----


class TestRunTaskFunction:
    async def test_run_task_convenience_mock(self) -> None:
        result = await run_task("score", "test", config=mock_config())
        assert result["score"] == 75