"""
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

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

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

        The classifier inspects *prompt* (which contains the formatted
        conversation text when called from ``UnifiedExtractor``) and decides
        whether the fast or smart client should handle this request.
        """
        result: ClassificationResult = self.classifier.classify(prompt)

        # 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 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()
