"""
Issue #124 (ENH-007): Tiered LLM client routing for extraction.

Generic wrapper that works with *any* two LLM clients (OpenAI, Anthropic,
Bedrock, etc.).  A TurnComplexityClassifier decides per-call whether to use
the fast (cheaper) or smart (more capable) client.  Satisfies the LLMClient
protocol, so it is a drop-in replacement for UnifiedExtractor.
"""

import logging
from dataclasses import dataclass, field
from typing import Optional

from .complexity_classifier import (
    TurnComplexityClassifier,
    ComplexityLevel,
    ClassificationResult,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Default model names per provider (used by MCAL.__init__ when the caller
# does not supply explicit fast_model / smart_model values).
# ---------------------------------------------------------------------------
DEFAULT_TIERED_MODELS: dict[str, dict[str, str]] = {
    "openai": {
        "fast": "gpt-4o-mini",
        "smart": "gpt-4o",
    },
    "anthropic": {
        "fast": "claude-3-5-haiku-20241022",
        "smart": "claude-sonnet-4-20250514",
    },
    "bedrock": {
        "fast": "llama-3.1-8b",
        "smart": "llama-3.3-70b",
    },
}


@dataclass
class RoutingMetrics:
    """Accumulated metrics for tiered model routing."""

    fast_calls: int = 0
    smart_calls: int = 0
    classifications: list = field(default_factory=list)

    @property
    def total_calls(self) -> int:
        return self.fast_calls + self.smart_calls

    @property
    def fast_ratio(self) -> float:
        return self.fast_calls / self.total_calls if self.total_calls > 0 else 0.0

    def to_dict(self) -> dict:
        return {
            "fast_calls": self.fast_calls,
            "smart_calls": self.smart_calls,
            "total_calls": self.total_calls,
            "fast_ratio": round(self.fast_ratio, 3),
            "recent_classifications": self.classifications[-5:],
        }

    def reset(self):
        self.fast_calls = 0
        self.smart_calls = 0
        self.classifications.clear()


class TieredLLMClient:
    """
    Generic tiered LLM wrapper that routes to fast / smart clients.

    Satisfies the ``LLMClient`` protocol::

        async def complete(self, prompt: str, system: Optional[str] = None) -> str

    The internal ``TurnComplexityClassifier`` inspects the prompt text and
    routes SIMPLE prompts to ``fast_client`` and COMPLEX prompts to
    ``smart_client``.

    Usage::

        fast = OpenAIClient(api_key=key, model="gpt-4o-mini")
        smart = OpenAIClient(api_key=key, model="gpt-4o")
        tiered = TieredLLMClient(fast, smart)

        # UnifiedExtractor sees a normal LLMClient:
        extractor = UnifiedExtractor(tiered)
    """

    def __init__(
        self,
        fast_client,
        smart_client,
        classifier: Optional[TurnComplexityClassifier] = None,
        max_history: int = 100,
    ):
        """
        Args:
            fast_client: LLMClient for simple / routine turns.
            smart_client: LLMClient for complex reasoning / decisions.
            classifier: Custom classifier (defaults to ``TurnComplexityClassifier()``).
            max_history: How many recent classification records to retain.
        """
        self.fast_client = fast_client
        self.smart_client = smart_client
        self.classifier = classifier or TurnComplexityClassifier()
        self.metrics = RoutingMetrics()
        self._max_history = max_history
        self._pending_classify_text: Optional[str] = None  # Issue #133

    # ------------------------------------------------------------------
    # LLMClient protocol
    # ------------------------------------------------------------------

    async def complete(self, prompt: str, system: Optional[str] = None) -> str:
        """
        Route completion to the appropriate model tier.

        The classifier inspects either the *classify_text* hint (if set
        via :meth:`set_classify_text` before the call) or *prompt* itself
        to decide whether the fast or smart client should handle this
        request.  Using ``set_classify_text`` lets callers (e.g.
        ``UnifiedExtractor``) supply the raw conversation text so the
        classifier sees user messages — not the extraction template
        boilerplate which always looks "complex" (Issue #133).
        """
        # Issue #133: prefer the raw-conversation hint if the caller set one.
        text_to_classify = self._pending_classify_text or prompt
        self._pending_classify_text = None  # consume it
        result: ClassificationResult = self.classifier.classify(text_to_classify)

        # Store classification for observability
        if len(self.metrics.classifications) >= self._max_history:
            self.metrics.classifications.pop(0)
        self.metrics.classifications.append(
            {
                "level": result.level.value,
                "confidence": round(result.confidence, 3),
                "reason": result.reason,
            }
        )

        if result.level == ComplexityLevel.COMPLEX:
            self.metrics.smart_calls += 1
            logger.info("Tiered routing → SMART model: %s", result.reason)
            return await self.smart_client.complete(prompt, system=system)
        else:
            self.metrics.fast_calls += 1
            logger.info("Tiered routing → FAST model: %s", result.reason)
            return await self.fast_client.complete(prompt, system=system)

    # ------------------------------------------------------------------
    # Metrics / observability
    # ------------------------------------------------------------------

    def set_classify_text(self, text: str) -> None:
        """
        Set the raw conversation text for the *next* ``complete()`` call.

        Issue #133: The extraction prompt template contains boilerplate that
        inflates the complexity score.  Callers should set the raw user
        message text here so the classifier sees it instead of the full
        prompt.  The value is automatically consumed (cleared) by the next
        ``complete()`` invocation.
        """
        self._pending_classify_text = text

    def get_routing_metrics(self) -> dict:
        """Return routing metrics dictionary."""
        return self.metrics.to_dict()

    def get_token_usage(self) -> dict:
        """Aggregate token usage from both underlying clients (if supported)."""
        usage: dict = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }
        for label, client in [("fast", self.fast_client), ("smart", self.smart_client)]:
            if hasattr(client, "get_token_usage"):
                cu = client.get_token_usage()
                usage["prompt_tokens"] += cu.get("prompt_tokens", 0)
                usage["completion_tokens"] += cu.get("completion_tokens", 0)
                usage["total_tokens"] += cu.get("total_tokens", 0)
                usage[f"{label}_tokens"] = cu.get("total_tokens", 0)
        return usage

    def reset_token_usage(self):
        """Reset token counters on both underlying clients (if supported)."""
        for client in [self.fast_client, self.smart_client]:
            if hasattr(client, "reset_token_usage"):
                client.reset_token_usage()
        self.metrics.reset()
