"""Model routing — automatic complexity-based model selection.

Routes simple messages (greetings, acknowledgments) to a cheaper/faster model
while keeping complex messages on the primary model. Saves ~60% cost on
simple interactions with zero quality loss.
"""

from __future__ import annotations

import logging
import re
from collections.abc import AsyncIterator
from dataclasses import dataclass

from qanot.providers.base import LLMProvider, ProviderResponse, StreamEvent

logger = logging.getLogger(__name__)

# --- Heuristic classifier ---------------------------------------------------

# Greeting / acknowledgment patterns (Uzbek, Russian, English)
_SIMPLE_PATTERNS: list[re.Pattern[str]] = [
    # Uzbek greetings & acknowledgments
    re.compile(r"^(salom|assalomu alaykum|salom aleykum|xayrli tong|xayrli kun|xayrli kech|rahmat|raxmat|yaxshi|ok|ha|yo['ʻ]q|keling|bo['ʻ]ldi|tushunarli|gap yo['ʻ]q|zo['ʻ]r|ajoyib|barakalla)\s*[.!?]*$", re.IGNORECASE),
    # Russian greetings & acknowledgments
    re.compile(r"^(привет|здравствуйте|добрый день|доброе утро|добрый вечер|спасибо|хорошо|ок|да|нет|понятно|ладно|пока|до свидания)\s*[.!?]*$", re.IGNORECASE),
    # English greetings & acknowledgments
    re.compile(r"^(hi|hello|hey|thanks|thank you|ok|okay|yes|no|sure|bye|goodbye|good morning|good evening|good night|got it|understood|cool|nice|great|awesome)\s*[.!?]*$", re.IGNORECASE),
]

# Indicators of complex messages
_COMPLEX_INDICATORS = [
    r"```",              # code blocks
    r"http[s]?://",      # URLs
    r"\d{3,}",           # long numbers (IDs, amounts)
    r"[{}\[\]]",         # JSON/data structures
    r"\b(explain|analyze|implement|create|build|write|fix|debug|compare|design)\b",
]
_COMPLEX_RE = re.compile("|".join(_COMPLEX_INDICATORS), re.IGNORECASE)


def classify_complexity(message: str) -> float:
    """Score message complexity from 0.0 (trivial) to 1.0 (complex).

    Uses fast heuristics — no LLM call needed.

    Scoring:
        0.0 - 0.2: Simple greetings, short acknowledgments
        0.2 - 0.5: Short questions, casual chat
        0.5 - 0.8: Medium complexity, some technical content
        0.8 - 1.0: Long/technical messages, code, multi-step requests
    """
    text = message.strip()

    if not text:
        return 0.0

    # Check for simple greeting/acknowledgment patterns
    for pattern in _SIMPLE_PATTERNS:
        if pattern.match(text):
            return 0.05

    score = 0.0

    # Length-based scoring (normalized to 0.0-0.4)
    length = len(text)
    if length < 20:
        score += 0.1
    elif length < 50:
        score += 0.2
    elif length < 150:
        score += 0.3
    else:
        score += 0.4

    # Word count factor
    words = text.split()
    word_count = len(words)
    if word_count > 30:
        score += 0.2
    elif word_count > 15:
        score += 0.1

    # Complex content indicators
    if _COMPLEX_RE.search(text):
        score += 0.3

    # Question marks (questions are slightly more complex)
    if "?" in text:
        score += 0.05

    # Multiple sentences
    sentence_count = len(re.split(r"[.!?]+", text))
    if sentence_count > 3:
        score += 0.1

    # Newlines (multi-line = more complex)
    if "\n" in text:
        score += 0.15

    return min(score, 1.0)


# --- Routing provider -------------------------------------------------------

@dataclass
class RoutingStats:
    """Track routing decisions for observability."""
    total: int = 0
    routed_cheap: int = 0
    routed_primary: int = 0

    @property
    def savings_pct(self) -> float:
        if self.total == 0:
            return 0.0
        return (self.routed_cheap / self.total) * 100


class RoutingProvider(LLMProvider):
    """3-tier model routing: Haiku → Sonnet → Opus based on complexity.

    Routes messages to the cheapest model that can handle them:
    - Haiku:  greetings, acknowledgments, simple questions (< 0.2)
    - Sonnet: general conversation, moderate tasks (0.2 - 0.6)
    - Opus:   complex tasks, tool calling, multi-step (> 0.6)

    Context-aware: if previous turn used Opus (tool calling, complex task),
    continuation messages ("ha", "davom et") stay on Opus.

    Cost savings: ~50-60% vs always using Opus.
    """

    def __init__(
        self,
        provider: LLMProvider,
        cheap_model: str = "claude-haiku-4-5-20251001",
        mid_model: str = "claude-sonnet-4-6",
        threshold: float = 0.3,
    ):
        self._provider = provider
        self._cheap_model = cheap_model
        self._mid_model = mid_model
        self._primary_model = provider.model  # Opus
        self._threshold = threshold
        self.model = provider.model
        self.stats = RoutingStats()
        # Track which model was used in the previous turn
        self._last_model: str = ""

    def _select_model(self, messages: list[dict]) -> str:
        """Pick model based on complexity + context continuity.

        Rules:
        1. If previous turn used Opus and user is continuing → stay on Opus
        2. Simple greetings/acks (score < 0.15) → Haiku
        3. Moderate messages (score < 0.5) → Sonnet
        4. Complex/tool-heavy (score >= 0.5) → Opus
        """
        user_text = ""
        for msg in reversed(messages):
            if msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, str):
                    user_text = content
                elif isinstance(content, list):
                    user_text = " ".join(
                        block.get("text", "")
                        for block in content
                        if isinstance(block, dict) and block.get("type") == "text"
                    )
                break

        msg_score = classify_complexity(user_text)
        ctx_score = self._assess_context(messages)
        effective = max(msg_score, ctx_score)

        self.stats.total += 1

        # Routing logic:
        # msg_score = what the USER is asking (their message complexity)
        # ctx_score = what was HAPPENING before (tool use, long responses)
        #
        # Key insight: if user sends a SIMPLE message, route to cheap model
        # REGARDLESS of context. Context only matters for ambiguous messages.
        # "salom" is always simple. "ha" after tool use is continuation.

        if msg_score < 0.1 and ctx_score < 0.5:
            # Pure greeting in calm context → Haiku
            self.stats.routed_cheap += 1
            selected = self._cheap_model
            logger.info("Routing → %s (greeting: msg=%.2f)", selected, msg_score)
        elif msg_score < 0.1 and ctx_score >= 0.5:
            # Short reply in active context ("ha", "yo'q" after tool use)
            # → stay on previous model (continuation)
            selected = self._last_model or self._mid_model
            self.stats.routed_primary += 1
            logger.info("Routing → %s (continuation: msg=%.2f, ctx=%.2f)", selected, msg_score, ctx_score)
        elif msg_score < 0.4:
            # Moderate message → Sonnet
            self.stats.routed_primary += 1
            selected = self._mid_model
            logger.info("Routing → %s (moderate: msg=%.2f)", selected, msg_score)
        else:
            # Complex message → Opus
            self.stats.routed_primary += 1
            selected = self._primary_model
            logger.info("Routing → %s (complex: msg=%.2f)", selected, msg_score)

        self._last_model = selected
        return selected

    @staticmethod
    def _assess_context(messages: list[dict]) -> float:
        """Score conversation context complexity (0.0 = fresh/simple, 1.0 = deep/complex).

        Checks:
        - Conversation depth (many turns = complex context)
        - Tool use in recent assistant messages (tool_use = complex task)
        - Previous assistant response length (long response = complex topic)
        """
        # Only look at the LAST 2 messages (immediate context, not history)
        recent = messages[-2:] if len(messages) >= 2 else messages
        if not recent:
            return 0.0

        score = 0.0

        for msg in recent:
            content = msg.get("content", "")

            # Tool use → active complex task
            if isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") in ("tool_use", "tool_result"):
                        score += 0.5
                        break

            # Long response → complex topic
            if isinstance(content, str) and len(content) > 500:
                score += 0.2
            elif isinstance(content, list):
                text_len = sum(
                    len(b.get("text", "")) for b in content
                    if isinstance(b, dict) and b.get("type") == "text"
                )
                if text_len > 500:
                    score += 0.2

        return min(score, 1.0)

    def _swap_model(self, model: str) -> str:
        """Swap the underlying provider's model, return the previous model."""
        prev = self._provider.model
        self._provider.model = model
        self.model = model
        return prev

    async def chat(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        system: str | None = None,
    ) -> ProviderResponse:
        """Route to appropriate model and call chat."""
        selected = self._select_model(messages)
        prev = self._swap_model(selected)
        try:
            return await self._provider.chat(messages, tools, system)
        finally:
            self._swap_model(prev)

    async def chat_stream(
        self,
        messages: list[dict],
        tools: list[dict] | None = None,
        system: str | None = None,
    ) -> AsyncIterator[StreamEvent]:
        """Route to appropriate model and stream."""
        selected = self._select_model(messages)
        prev = self._swap_model(selected)
        try:
            async for event in self._provider.chat_stream(messages, tools, system):
                yield event
        finally:
            self._swap_model(prev)

    def status(self) -> dict:
        """Return routing statistics."""
        return {
            "primary_model": self._primary_model,
            "cheap_model": self._cheap_model,
            "threshold": self._threshold,
            "stats": {
                "total": self.stats.total,
                "routed_cheap": self.stats.routed_cheap,
                "routed_primary": self.stats.routed_primary,
                "savings_pct": round(self.stats.savings_pct, 1),
            },
        }
