from __future__ import annotations

import litellm

from costrouter.exceptions import ScreenerError
from costrouter.models import CostTier
from costrouter.pricing import calculate_cost
from costrouter.router.base import BaseRouter

CLASSIFIER_SYSTEM_PROMPT = """You are a query complexity classifier for an AI routing system.
Analyze the user's query and classify its complexity into exactly one category.

Categories:
- SIMPLE: Basic facts, greetings, simple math, translations, yes/no questions
- MODERATE: Summaries, concept explanations, moderate analysis, short writing tasks
- COMPLEX: Multi-step reasoning, code generation, creative writing, detailed analysis
- EXPERT: Advanced research, complex proofs, specialized domain expertise, system design

Respond with ONLY the category name (SIMPLE, MODERATE, COMPLEX, or EXPERT) and nothing else."""

_TIER_MAP: dict[str, CostTier] = {
    "SIMPLE": CostTier.SIMPLE,
    "MODERATE": CostTier.MODERATE,
    "COMPLEX": CostTier.COMPLEX,
    "EXPERT": CostTier.EXPERT,
}

_LLM_CONFIDENCE = 0.85


def _parse_tier(text: str) -> CostTier | None:
    """Parse the classifier response into a CostTier."""
    cleaned = text.strip().upper()
    return _TIER_MAP.get(cleaned)


class LLMRouter(BaseRouter):
    """Routes using a cheap LLM to classify complexity."""

    def __init__(
        self,
        screener_model: str = "gpt-4o-mini",
        api_key: str | None = None,
    ) -> None:
        self.screener_model = screener_model
        self.api_key = api_key

    def _build_messages(self, messages: list[dict]) -> list[dict]:
        """Build the classification prompt from the full conversation."""
        return [
            {"role": "system", "content": CLASSIFIER_SYSTEM_PROMPT},
            *messages,
        ]

    def classify(self, messages: list[dict]) -> tuple[CostTier, float, float]:
        kwargs: dict = {
            "model": self.screener_model,
            "messages": self._build_messages(messages),
            "max_tokens": 10,
            "temperature": 0,
        }
        if self.api_key:
            kwargs["api_key"] = self.api_key

        try:
            response = litellm.completion(**kwargs)
        except Exception as exc:
            raise ScreenerError(f"LLM screening failed: {exc}") from exc

        screener_cost = calculate_cost(response)
        content = response.choices[0].message.content
        tier = _parse_tier(content)
        if tier is None:
            return CostTier.MODERATE, 0.5, screener_cost

        return tier, _LLM_CONFIDENCE, screener_cost

    async def aclassify(self, messages: list[dict]) -> tuple[CostTier, float, float]:
        kwargs: dict = {
            "model": self.screener_model,
            "messages": self._build_messages(messages),
            "max_tokens": 10,
            "temperature": 0,
        }
        if self.api_key:
            kwargs["api_key"] = self.api_key

        try:
            response = await litellm.acompletion(**kwargs)
        except Exception as exc:
            raise ScreenerError(f"LLM screening failed: {exc}") from exc

        screener_cost = calculate_cost(response)
        content = response.choices[0].message.content
        tier = _parse_tier(content)
        if tier is None:
            return CostTier.MODERATE, 0.5, screener_cost

        return tier, _LLM_CONFIDENCE, screener_cost
