"""Tests for the smart model router."""

from __future__ import annotations

from agent_scout.llm.router import (
    ModelRouter,
    RouterConfig,
    RoutingRule,
    TaskCategory,
    TaskComplexity,
    classify_task,
)

# --- classify_task tests ---


def test_classify_lookup():
    cat, comp = classify_task("Find the cheapest flight to Paris")
    assert cat == TaskCategory.LOOKUP
    assert comp == TaskComplexity.SIMPLE


def test_classify_summarization():
    cat, comp = classify_task("Summarize the search results for me")
    assert cat == TaskCategory.SUMMARIZATION
    assert comp == TaskComplexity.MODERATE


def test_classify_reasoning():
    cat, comp = classify_task("Analyze the pros and cons of each option")
    assert cat == TaskCategory.REASONING
    assert comp == TaskComplexity.COMPLEX


def test_classify_coding():
    cat, comp = classify_task("Write a function to parse CSV files")
    assert cat == TaskCategory.CODING
    assert comp == TaskComplexity.COMPLEX


def test_classify_conversation_default():
    cat, comp = classify_task("Hello, how are you today?")
    assert cat == TaskCategory.CONVERSATION
    assert comp == TaskComplexity.MODERATE


def test_classify_tool_use_when_tools_available():
    cat, comp = classify_task("Hello, how are you?", tool_count=3)
    assert cat == TaskCategory.TOOL_USE
    assert comp == TaskComplexity.MODERATE


def test_classify_bumps_complexity_at_many_steps():
    cat, comp = classify_task("Search for flights", step_count=5)
    assert comp == TaskComplexity.COMPLEX


# --- ModelRouter tests ---


def test_router_simple_task():
    router = ModelRouter()
    model = router.route("Find the cheapest flight")
    assert model == router.config.fast_model


def test_router_complex_task():
    router = ModelRouter()
    model = router.route("Analyze and compare these three options")
    assert model == router.config.best_model


def test_router_moderate_task():
    router = ModelRouter()
    model = router.route("Summarize the article")
    assert model == router.config.balanced_model


def test_router_override_model():
    router = ModelRouter()
    model = router.route("Find flights", override_model="openai/gpt-4o")
    assert model == "openai/gpt-4o"


def test_router_disabled():
    config = RouterConfig(enabled=False)
    router = ModelRouter(config)
    model = router.route("Analyze complex data")
    assert model == config.balanced_model


def test_router_custom_rule():
    rule = RoutingRule(
        category=TaskCategory.LOOKUP,
        complexity=TaskComplexity.SIMPLE,
        model="groq/llama-3-8b",
    )
    config = RouterConfig(rules=[rule])
    router = ModelRouter(config)
    model = router.route("Find the weather in NYC")
    assert model == "groq/llama-3-8b"


def test_router_custom_rule_no_match():
    rule = RoutingRule(
        category=TaskCategory.CODING,
        complexity=TaskComplexity.COMPLEX,
        model="openai/gpt-4",
    )
    config = RouterConfig(rules=[rule])
    router = ModelRouter(config)
    # This is a lookup, not coding — rule shouldn't match
    model = router.route("Find the weather in NYC")
    assert model == config.fast_model


def test_router_custom_models():
    config = RouterConfig(
        fast_model="groq/llama-3-8b",
        balanced_model="openai/gpt-4o-mini",
        best_model="anthropic/claude-opus-4-20250514",
    )
    router = ModelRouter(config)

    fast = router.route("Search for the price")
    assert fast == "groq/llama-3-8b"

    best = router.route("Analyze and recommend the best option")
    assert best == "anthropic/claude-opus-4-20250514"


def test_routing_rule_matches():
    rule = RoutingRule(category=TaskCategory.LOOKUP)
    assert rule.matches(TaskCategory.LOOKUP, TaskComplexity.SIMPLE)
    assert not rule.matches(TaskCategory.CODING, TaskComplexity.SIMPLE)

    # Rule with no filters matches everything
    rule_any = RoutingRule()
    assert rule_any.matches(TaskCategory.CODING, TaskComplexity.COMPLEX)
