"""Tests for model registry and routing strategy."""

import pytest

from proxym.providers import (
    MODELS,
    Capability,
    ModelSpec,
    TaskTier,
    cheapest_model,
    get_model,
    models_for_tier,
)
from proxym.router import AnalysisResult, analyze_request
from proxym.router.strategy import CostLedger, Router


class TestModelRegistry:
    """Verify the model registry is correctly populated."""

    def test_has_15_models(self):
        assert len(MODELS) == 15

    def test_all_models_have_provider(self):
        for mid, m in MODELS.items():
            assert m.provider, f"{mid} missing provider"

    def test_all_models_have_litellm_string(self):
        for mid, m in MODELS.items():
            assert m.litellm_model, f"{mid} missing litellm_model"

    def test_opus_is_most_expensive(self):
        opus = MODELS["anthropic/opus-4.6"]
        for mid, m in MODELS.items():
            if mid == "anthropic/opus-4.6":
                continue
            # Opus should have the highest output cost among Anthropic models
            if m.provider == "anthropic":
                assert opus.output_cost >= m.output_cost

    def test_local_models_are_free(self):
        for mid, m in MODELS.items():
            if m.provider == "ollama":
                assert m.input_cost == 0.0
                assert m.output_cost == 0.0

    def test_get_model_by_id(self):
        m = get_model("anthropic/sonnet-4.6")
        assert m is not None
        assert m.display_name == "Claude Sonnet 4.6"

    def test_get_model_missing(self):
        assert get_model("nonexistent/model") is None


class TestModelsForTier:
    """Verify tier-based model filtering."""

    def test_trivial_returns_cheap_models(self):
        models = models_for_tier(TaskTier.TRIVIAL)
        assert len(models) > 0
        # First model should be one of the cheapest
        assert models[0].effective_cost <= 5.0

    def test_deep_reasoning_includes_opus(self):
        models = models_for_tier(TaskTier.DEEP_REASONING)
        model_ids = {m.id for m in models}
        assert "anthropic/opus-4.6" in model_ids

    def test_trivial_excludes_opus(self):
        models = models_for_tier(TaskTier.TRIVIAL)
        model_ids = {m.id for m in models}
        assert "anthropic/opus-4.6" not in model_ids

    def test_filter_by_capability(self):
        models = models_for_tier(
            TaskTier.STANDARD,
            required_caps=frozenset({Capability.THINKING}),
        )
        for m in models:
            assert Capability.THINKING in m.capabilities

    def test_filter_by_provider(self):
        models = models_for_tier(
            TaskTier.STANDARD,
            available_providers={"anthropic"},
        )
        for m in models:
            assert m.provider == "anthropic"

    def test_sorted_by_cost(self):
        models = models_for_tier(TaskTier.STANDARD)
        costs = [m.effective_cost for m in models]
        assert costs == sorted(costs)

    def test_cheapest_model_returns_cheapest(self):
        m = cheapest_model(TaskTier.TRIVIAL)
        assert m is not None
        all_trivial = models_for_tier(TaskTier.TRIVIAL)
        assert m.effective_cost == all_trivial[0].effective_cost


class TestCostLedger:
    """Verify budget tracking and enforcement."""

    def test_initial_can_spend(self):
        ledger = CostLedger(daily_limit=5.0, monthly_limit=60.0)
        assert ledger.can_spend(1.0) is True

    def test_daily_limit_enforcement(self):
        ledger = CostLedger(daily_limit=1.0, monthly_limit=60.0)
        ledger.record(0.8)
        assert ledger.can_spend(0.3) is False  # 0.8 + 0.3 > 1.0
        assert ledger.can_spend(0.1) is True   # 0.8 + 0.1 < 1.0

    def test_monthly_limit_enforcement(self):
        ledger = CostLedger(daily_limit=100.0, monthly_limit=2.0)
        ledger.record(1.5)
        assert ledger.can_spend(0.6) is False  # 1.5 + 0.6 > 2.0
        assert ledger.can_spend(0.4) is True

    def test_zero_limit_means_unlimited(self):
        ledger = CostLedger(daily_limit=0, monthly_limit=0)
        ledger.record(9999.0)
        assert ledger.can_spend(9999.0) is True

    def test_remaining_decreases(self):
        ledger = CostLedger(daily_limit=5.0, monthly_limit=60.0)
        assert ledger.daily_remaining == 5.0
        ledger.record(2.0)
        assert ledger.daily_remaining == 3.0
        assert ledger.monthly_remaining == 58.0


class TestRouter:
    """Verify the router picks the right model for each scenario."""

    def test_trivial_task_picks_cheap_model(self, router):
        analysis = AnalysisResult(
            tier=TaskTier.TRIVIAL,
            capabilities=frozenset(),
        )
        decision = router.route(analysis)
        # Should pick one of the cheapest models
        assert decision.model.effective_cost < 5.0
        assert decision.estimated_cost_usd < 0.1

    def test_complex_task_picks_capable_model(self, router):
        analysis = AnalysisResult(
            tier=TaskTier.COMPLEX,
            capabilities=frozenset({Capability.CODE, Capability.REASONING}),
            estimated_output_tokens=5000,
        )
        decision = router.route(analysis)
        assert Capability.CODE in decision.model.capabilities

    def test_deep_reasoning_can_pick_opus(self, router):
        # Override max_request_cost to allow Opus
        router.max_request_cost = 10.0
        analysis = AnalysisResult(
            tier=TaskTier.DEEP_REASONING,
            capabilities=frozenset({Capability.THINKING, Capability.CODE}),
            estimated_output_tokens=6000,
        )
        decision = router.route(analysis)
        assert Capability.THINKING in decision.model.capabilities

    def test_explicit_model_passthrough(self, router):
        analysis = AnalysisResult(tier=TaskTier.STANDARD, capabilities=frozenset())
        decision = router.route(analysis, explicit_model="anthropic/opus-4.6")
        assert decision.model.id == "anthropic/opus-4.6"

    def test_explicit_model_partial_match(self, router):
        analysis = AnalysisResult(tier=TaskTier.STANDARD, capabilities=frozenset())
        decision = router.route(analysis, explicit_model="opus-4.6")
        assert decision.model.id == "anthropic/opus-4.6"

    def test_fallbacks_are_populated(self, router):
        analysis = AnalysisResult(
            tier=TaskTier.STANDARD,
            capabilities=frozenset(),  # no specific caps → more candidates
        )
        decision = router.route(analysis)
        # With no capability filter, should have multiple candidates
        assert decision.model is not None

    def test_budget_exceeded_falls_to_free(self, cheap_router):
        # Spend almost entire budget
        cheap_router.ledger.record(0.95)
        analysis = AnalysisResult(
            tier=TaskTier.STANDARD,
            capabilities=frozenset({Capability.CODE}),
            estimated_output_tokens=5000,
        )
        decision = cheap_router.route(analysis)
        # Should fall to something cheap/free
        assert decision.model.effective_cost <= 2.0

    def test_latency_tracking(self, router):
        router.record_latency("anthropic/sonnet-4.6", 1.5)
        router.record_latency("anthropic/sonnet-4.6", 2.0)
        ema = router._latency_ema["anthropic/sonnet-4.6"]
        assert 1.0 < ema < 2.5

    def test_routing_time_is_fast(self, router):
        analysis = AnalysisResult(
            tier=TaskTier.STANDARD,
            capabilities=frozenset({Capability.CODE}),
        )
        decision = router.route(analysis)
        assert decision.routing_time_ms < 50  # should be < 1ms typically

    def test_vision_requires_capable_model(self, router):
        router.max_request_cost = 10.0
        analysis = AnalysisResult(
            tier=TaskTier.DEEP_REASONING,  # use deep tier where vision models exist
            capabilities=frozenset({Capability.VISION}),
        )
        decision = router.route(analysis)
        assert Capability.VISION in decision.model.capabilities


class TestAnalyzeRequestHeuristics:
    """Verify request analysis heuristics after refactor."""

    def test_long_message_bumps_to_complex_tier(self):
        """Long messages (many tokens) bump to complex tier."""
        # Create a long message (simulating many tokens)
        long_content = "Explain quantum computing. " * 100  # ~400 tokens
        result = analyze_request(long_content)
        # Should be at least STANDARD, possibly COMPLEX
        assert result.tier in (TaskTier.STANDARD, TaskTier.COMPLEX, TaskTier.DEEP_REASONING)

    def test_image_block_adds_vision_capability(self):
        """Image blocks in request add VISION capability."""
        messages = [
            {"role": "user", "content": [
                {"type": "text", "text": "What's in this image?"},
                {"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
            ]},
        ]
        result = analyze_request(messages)
        assert Capability.VISION in result.capabilities

    def test_code_keywords_bump_tier(self):
        """Code-related keywords bump tier to at least STANDARD."""
        code_prompt = "Write a Python function to sort a list using quicksort"
        result = analyze_request(code_prompt)
        assert result.tier in (TaskTier.STANDARD, TaskTier.COMPLEX)
        assert Capability.CODE in result.capabilities

    def test_many_files_in_context_bumps_tier(self):
        """Many files in context bump tier."""
        # Simulate request with many files in context
        files = [f"file{i}.py" for i in range(50)]
        result = analyze_request("Review these files", context_files=files)
        # Large context should bump to at least STANDARD
        assert result.tier in (TaskTier.STANDARD, TaskTier.COMPLEX, TaskTier.DEEP_REASONING)

    def test_explicit_tier_header_overrides_heuristics(self):
        """X-Task-Tier header overrides automatic heuristic detection."""
        # Simple content would normally be TRIVIAL
        simple_content = "Hi"
        result = analyze_request(simple_content, explicit_tier=TaskTier.DEEP_REASONING)
        assert result.tier == TaskTier.DEEP_REASONING

    def test_simple_greeting_is_trivial(self):
        """Simple greeting should be TRIVIAL tier."""
        result = analyze_request("Hello")
        assert result.tier == TaskTier.TRIVIAL

    def test_reasoning_request_adds_thinking_capability(self):
        """Requests mentioning reasoning/thinking add THINKING capability."""
        reasoning_prompt = "Explain step by step how to solve this math problem"
        result = analyze_request(reasoning_prompt)
        assert Capability.THINKING in result.capabilities

    def test_tool_use_request_adds_tools_capability(self):
        """Requests mentioning tools/functions add TOOLS capability."""
        tool_prompt = "Use the calculator tool to compute 123 * 456"
        result = analyze_request(tool_prompt)
        assert Capability.TOOLS in result.capabilities

