from __future__ import annotations

import time
from datetime import datetime, timezone
from typing import Any

import litellm

from costrouter.analytics import Analytics, RoutingRecord
from costrouter.exceptions import AllModelsFailedError, NoModelForTierError, ScreenerError
from costrouter.models import (
    CostTier,
    ModelConfig,
    RoutedResponse,
    RoutingResult,
    TIER_ORDER,
)
from costrouter.pricing import calculate_cost, estimate_max_cost
from costrouter.router.base import BaseRouter
from costrouter.router.cascade import CascadeRouter
from costrouter.router.heuristic import HeuristicRouter
from costrouter.router.llm import LLMRouter


def _parse_models(models: list[str] | dict[str, str] | list[ModelConfig]) -> list[ModelConfig]:
    """Normalize model config into a list of ModelConfig.

    Accepts:
      - list[str]: ordered cheap→expensive, auto-assigns tiers by position
      - dict[str, str]: tier name → model name
      - list[ModelConfig]: full control
    """
    if isinstance(models, dict):
        return [
            ModelConfig(name=model_name, tier=CostTier(tier_name))
            for tier_name, model_name in models.items()
        ]

    if not models:
        return []

    # list[str] — auto-assign tiers by position
    if isinstance(models[0], str):
        str_models: list[str] = models  # type: ignore[assignment]
        n = len(str_models)
        if n == 1:
            tiers = [CostTier.MODERATE]
        elif n == 2:
            tiers = [CostTier.SIMPLE, CostTier.COMPLEX]
        elif n == 3:
            tiers = [CostTier.SIMPLE, CostTier.MODERATE, CostTier.COMPLEX]
        else:
            tiers = list(TIER_ORDER[:4])
            # If more than 4, extra models share the EXPERT tier
            for _ in range(n - 4):
                tiers.append(CostTier.EXPERT)
        return [
            ModelConfig(name=name, tier=tier)
            for name, tier in zip(str_models, tiers)
        ]

    return list(models)


def _find_model_for_tier(
    models: list[ModelConfig], tier: CostTier
) -> ModelConfig | None:
    """Find a model configured for the given tier."""
    for m in models:
        if m.tier == tier:
            return m
    return None


def _resolve_model(
    models: list[ModelConfig], tier: CostTier
) -> ModelConfig:
    """Find the best model for a tier, falling back to nearest available tier.

    Prefers cheaper (lower) tiers first to maximize cost savings.
    """
    # Exact match first
    model = _find_model_for_tier(models, tier)
    if model:
        return model

    tier_idx = TIER_ORDER.index(tier)

    # Search downward first (prefer cheaper)
    for i in range(tier_idx - 1, -1, -1):
        model = _find_model_for_tier(models, TIER_ORDER[i])
        if model:
            return model

    # Then search upward
    for i in range(tier_idx + 1, len(TIER_ORDER)):
        model = _find_model_for_tier(models, TIER_ORDER[i])
        if model:
            return model

    raise NoModelForTierError(f"No model available for tier {tier.value} or any adjacent tier")


def _get_most_expensive_model(models: list[ModelConfig]) -> str:
    """Return the name of the model in the highest tier."""
    for tier in reversed(TIER_ORDER):
        model = _find_model_for_tier(models, tier)
        if model:
            return model.name
    return models[0].name


_RESERVED_KEYS = frozenset({"model", "messages"})


def _build_call_kwargs(
    model_config: ModelConfig, messages: list[dict], **kwargs: Any
) -> dict:
    """Build kwargs for litellm.completion()."""
    call_kwargs: dict[str, Any] = {}

    # Apply extra_params first so they can't overwrite core fields
    if model_config.extra_params:
        for k, v in model_config.extra_params.items():
            if k not in _RESERVED_KEYS:
                call_kwargs[k] = v

    # User kwargs next (can override extra_params but not core fields)
    for k, v in kwargs.items():
        if k not in _RESERVED_KEYS:
            call_kwargs[k] = v

    # Core fields last — always authoritative
    call_kwargs["model"] = model_config.name
    call_kwargs["messages"] = messages

    if model_config.api_key:
        call_kwargs["api_key"] = model_config.api_key
    if model_config.api_base:
        call_kwargs["api_base"] = model_config.api_base

    return call_kwargs


class CostRouter:
    """Main CostRouter client — routes prompts to the optimal LLM by complexity."""

    def __init__(
        self,
        models: list[str] | dict[str, str] | list[ModelConfig],
        screener: str | None = None,
        strategy: str | BaseRouter = "cascade",
        screener_api_key: str | None = None,
        confidence_threshold: float = 0.7,
        default_tier: CostTier = CostTier.MODERATE,
        fallback: bool = True,
    ) -> None:
        self._models = _parse_models(models)
        if not self._models:
            raise ValueError("At least one model must be provided")

        # Default screener to the cheapest model
        if screener is None:
            screener = self._models[0].name

        self._default_tier = default_tier
        self._fallback = fallback
        self._analytics = Analytics()
        self._most_expensive = _get_most_expensive_model(self._models)

        self._router: BaseRouter
        if isinstance(strategy, BaseRouter):
            self._router = strategy
            self._strategy = "custom"
        elif strategy == "heuristic":
            self._router = HeuristicRouter()
            self._strategy = strategy
        elif strategy == "llm":
            self._router = LLMRouter(screener_model=screener, api_key=screener_api_key)
            self._strategy = strategy
        else:
            self._router = CascadeRouter(
                screener_model=screener,
                api_key=screener_api_key,
                confidence_threshold=confidence_threshold,
            )
            self._strategy = "cascade"

    @property
    def analytics(self) -> Analytics:
        return self._analytics

    def completion(self, messages: list[dict], **kwargs: Any) -> RoutedResponse:
        """Route and complete a chat request synchronously."""
        start = time.monotonic()

        # Step 1: Classify (fall back to default_tier on screener failure)
        try:
            tier, confidence, screener_cost = self._router.classify(messages)
        except ScreenerError:
            tier, confidence, screener_cost = self._default_tier, 0.0, 0.0

        # Step 2: Resolve model
        model_config = _resolve_model(self._models, tier)

        # Step 3: Call the model (with fallback)
        response, call_cost, model_used = self._call_with_fallback(
            model_config, tier, messages, **kwargs
        )

        elapsed_ms = (time.monotonic() - start) * 1000
        total_cost = screener_cost + call_cost
        max_cost = estimate_max_cost(
            self._most_expensive, response=response, messages=messages
        )
        saved = max(0.0, max_cost - total_cost)

        routing = RoutingResult(
            tier=tier,
            model=model_used,
            confidence=confidence,
            router_type=self._strategy,
            screener_cost=screener_cost,
        )

        prompt_length = sum(len(m.get("content", "")) for m in messages)
        self._analytics.record(
            RoutingRecord(
                timestamp=datetime.now(timezone.utc),
                prompt_length=prompt_length,
                tier=tier,
                model_used=model_used,
                router_type=self._strategy,
                screener_cost=screener_cost,
                call_cost=call_cost,
                total_cost=total_cost,
                max_cost_estimate=max_cost,
                latency_ms=elapsed_ms,
            )
        )

        return RoutedResponse(
            response=response,
            routing=routing,
            cost=call_cost,
            total_cost=total_cost,
            saved=saved,
        )

    async def acompletion(self, messages: list[dict], **kwargs: Any) -> RoutedResponse:
        """Route and complete a chat request asynchronously."""
        start = time.monotonic()

        # Step 1: Classify (fall back to default_tier on screener failure)
        try:
            tier, confidence, screener_cost = await self._router.aclassify(messages)
        except ScreenerError:
            tier, confidence, screener_cost = self._default_tier, 0.0, 0.0

        # Step 2: Resolve model
        model_config = _resolve_model(self._models, tier)

        # Step 3: Call the model (with fallback)
        response, call_cost, model_used = await self._acall_with_fallback(
            model_config, tier, messages, **kwargs
        )

        elapsed_ms = (time.monotonic() - start) * 1000
        total_cost = screener_cost + call_cost
        max_cost = estimate_max_cost(
            self._most_expensive, response=response, messages=messages
        )
        saved = max(0.0, max_cost - total_cost)

        routing = RoutingResult(
            tier=tier,
            model=model_used,
            confidence=confidence,
            router_type=self._strategy,
            screener_cost=screener_cost,
        )

        prompt_length = sum(len(m.get("content", "")) for m in messages)
        self._analytics.record(
            RoutingRecord(
                timestamp=datetime.now(timezone.utc),
                prompt_length=prompt_length,
                tier=tier,
                model_used=model_used,
                router_type=self._strategy,
                screener_cost=screener_cost,
                call_cost=call_cost,
                total_cost=total_cost,
                max_cost_estimate=max_cost,
                latency_ms=elapsed_ms,
            )
        )

        return RoutedResponse(
            response=response,
            routing=routing,
            cost=call_cost,
            total_cost=total_cost,
            saved=saved,
        )

    def _call_with_fallback(
        self,
        model_config: ModelConfig,
        tier: CostTier,
        messages: list[dict],
        **kwargs: Any,
    ) -> tuple[Any, float, str]:
        """Try the target model, then fall back to higher tiers on failure."""
        errors: list[Exception] = []
        tried: set[str] = set()

        # Build ordered list of models to try
        candidates = self._get_fallback_candidates(model_config, tier)

        for candidate in candidates:
            if candidate.name in tried:
                continue
            tried.add(candidate.name)

            try:
                call_kwargs = _build_call_kwargs(candidate, messages, **kwargs)
                response = litellm.completion(**call_kwargs)
                cost = calculate_cost(response)
                return response, cost, candidate.name
            except Exception as exc:
                errors.append(exc)
                if not self._fallback:
                    break

        raise AllModelsFailedError(
            f"All models failed. Errors: {errors}"
        )

    async def _acall_with_fallback(
        self,
        model_config: ModelConfig,
        tier: CostTier,
        messages: list[dict],
        **kwargs: Any,
    ) -> tuple[Any, float, str]:
        """Async version of _call_with_fallback."""
        errors: list[Exception] = []
        tried: set[str] = set()

        candidates = self._get_fallback_candidates(model_config, tier)

        for candidate in candidates:
            if candidate.name in tried:
                continue
            tried.add(candidate.name)

            try:
                call_kwargs = _build_call_kwargs(candidate, messages, **kwargs)
                response = await litellm.acompletion(**call_kwargs)
                cost = calculate_cost(response)
                return response, cost, candidate.name
            except Exception as exc:
                errors.append(exc)
                if not self._fallback:
                    break

        raise AllModelsFailedError(
            f"All models failed. Errors: {errors}"
        )

    def _get_fallback_candidates(
        self, primary: ModelConfig, tier: CostTier
    ) -> list[ModelConfig]:
        """Get ordered list of models to try: primary first, then higher tiers."""
        candidates = [primary]

        if self._fallback:
            tier_idx = TIER_ORDER.index(tier)
            for i in range(tier_idx + 1, len(TIER_ORDER)):
                model = _find_model_for_tier(self._models, TIER_ORDER[i])
                if model and model.name != primary.name:
                    candidates.append(model)

        return candidates
