from costrouter.models import CostTier
from costrouter.router.heuristic import HeuristicRouter


def _msgs(text: str) -> list[dict]:
    return [{"role": "user", "content": text}]


class TestHeuristicRouter:
    def setup_method(self):
        self.router = HeuristicRouter()

    def test_short_prompt_is_simple(self):
        tier, conf, cost = self.router.classify(_msgs("Hi"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.7
        assert cost == 0.0

    def test_greeting_hello(self):
        tier, conf, cost = self.router.classify(_msgs("Hello, how are you doing today?"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.9

    def test_greeting_thanks(self):
        tier, conf, _ = self.router.classify(_msgs("Thanks for the help with that!"))
        assert tier == CostTier.SIMPLE
        assert conf == 0.9

    def test_code_keyword_write_function(self):
        tier, conf, _ = self.router.classify(
            _msgs("Write a function to sort a list in Python")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.7

    def test_code_keyword_debug(self):
        tier, conf, _ = self.router.classify(
            _msgs("Debug this error in my application")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.7

    def test_code_block(self):
        tier, conf, _ = self.router.classify(
            _msgs("What does this do?\n```python\nprint('hello')\n```")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.7

    def test_expert_prove(self):
        tier, conf, _ = self.router.classify(
            _msgs("Prove that the square root of 2 is irrational using contradiction")
        )
        assert tier == CostTier.EXPERT
        assert conf == 0.6

    def test_expert_differential_equation(self):
        tier, conf, _ = self.router.classify(
            _msgs("Solve this differential equation: dy/dx = 3x^2 + 2x")
        )
        assert tier == CostTier.EXPERT
        assert conf == 0.6

    def test_creative_writing(self):
        tier, conf, _ = self.router.classify(
            _msgs("Write a story about a robot who learns to paint")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.6

    def test_simple_factual_what_is(self):
        tier, conf, _ = self.router.classify(
            _msgs("What is the capital of France?")
        )
        assert tier == CostTier.SIMPLE
        assert conf == 0.6

    def test_simple_factual_who_is(self):
        tier, conf, _ = self.router.classify(
            _msgs("Who is the president of the United States?")
        )
        assert tier == CostTier.SIMPLE
        assert conf == 0.6

    def test_analysis_compare(self):
        tier, conf, _ = self.router.classify(
            _msgs("Compare the economic policies of the two candidates in detail")
        )
        assert tier == CostTier.MODERATE
        assert conf == 0.5

    def test_analysis_explain_why(self):
        tier, conf, _ = self.router.classify(
            _msgs("Explain why the sky appears blue during the day")
        )
        assert tier == CostTier.MODERATE
        assert conf == 0.5

    def test_default_moderate(self):
        tier, conf, _ = self.router.classify(
            _msgs("Tell me something interesting about octopuses and their behavior")
        )
        assert tier == CostTier.MODERATE
        assert conf == 0.3

    def test_screener_cost_always_zero(self):
        _, _, cost = self.router.classify(_msgs("anything"))
        assert cost == 0.0

    def test_empty_messages_uses_default(self):
        tier, conf, _ = self.router.classify([])
        assert tier == CostTier.SIMPLE  # empty string < 20 chars
        assert conf == 0.7

    def test_system_message_only(self):
        tier, conf, _ = self.router.classify(
            [{"role": "system", "content": "You are helpful."}]
        )
        assert tier == CostTier.SIMPLE  # no user message -> empty -> short

    def test_greeting_with_code_keyword_prefers_complex(self):
        """'Hey, write a function...' should be COMPLEX, not SIMPLE."""
        tier, conf, _ = self.router.classify(
            _msgs("Hey, write a function to parse JSON in Python")
        )
        assert tier == CostTier.COMPLEX
        assert conf == 0.7

    def test_greeting_with_expert_keyword_prefers_expert(self):
        """'Hi, prove this theorem...' should be EXPERT, not SIMPLE."""
        tier, conf, _ = self.router.classify(
            _msgs("Hi, can you prove that there are infinitely many primes?")
        )
        assert tier == CostTier.EXPERT
        assert conf == 0.6

    def test_multi_turn_uses_last_user(self):
        messages = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "user", "content": "Compare the GDP of Japan and Germany in detail please"},
        ]
        tier, _, _ = self.router.classify(messages)
        assert tier == CostTier.MODERATE  # "compare" triggers analysis
