"""Routing strategy: maps AnalysisResult → ModelSpec, respecting budget and availability.

The router maintains a cost ledger (Redis-backed) and refuses requests that would
exceed daily/monthly budgets.  It also tracks per-model latency to prefer faster
deployments when quality is equivalent.
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field

import structlog

from proxym.providers import (
    Capability,
    ModelSpec,
    TaskTier,
    cheapest_model,
    get_model,
    models_for_tier,
)
from proxym.observability import log_call
from proxym.router import AnalysisResult
from proxym.config import get_settings

logger = structlog.get_logger()


@dataclass
class RoutingDecision:
    """Complete routing decision with audit trail."""

    model: ModelSpec
    fallbacks: list[ModelSpec] = field(default_factory=list)
    estimated_cost_usd: float = 0.0
    analysis: AnalysisResult | None = None
    routing_time_ms: float = 0.0
    reason: str = ""


class CostLedger:
    """In-memory cost ledger with Redis persistence.

    Tracks spend per day and month.  Call ``record()`` after each completion.
    """

    def __init__(self, daily_limit: float = 5.0, monthly_limit: float = 60.0):
        self.daily_limit = daily_limit
        self.monthly_limit = monthly_limit
        self._daily_spend: float = 0.0
        self._monthly_spend: float = 0.0
        self._day_key: str = ""
        self._month_key: str = ""
        self._rotate_keys()

    def _rotate_keys(self) -> None:
        today = time.strftime("%Y-%m-%d")
        month = time.strftime("%Y-%m")
        if today != self._day_key:
            self._day_key = today
            self._daily_spend = 0.0
        if month != self._month_key:
            self._month_key = month
            self._monthly_spend = 0.0

    def can_spend(self, estimated_usd: float) -> bool:
        self._rotate_keys()
        if self.daily_limit > 0 and self._daily_spend + estimated_usd > self.daily_limit:
            return False
        if self.monthly_limit > 0 and self._monthly_spend + estimated_usd > self.monthly_limit:
            return False
        return True

    def record(self, actual_usd: float) -> None:
        self._rotate_keys()
        self._daily_spend += actual_usd
        self._monthly_spend += actual_usd

    @property
    def daily_remaining(self) -> float:
        self._rotate_keys()
        return max(0, self.daily_limit - self._daily_spend)

    @property
    def monthly_remaining(self) -> float:
        self._rotate_keys()
        return max(0, self.monthly_limit - self._monthly_spend)


class Router:
    """Stateful router that selects the optimal model per request."""

    def __init__(
        self,
        available_providers: set[str],
        ledger: CostLedger | None = None,
        max_request_cost: float = 2.0,
    ):
        self.available_providers = available_providers
        self.ledger = ledger or CostLedger()
        self.max_request_cost = max_request_cost
        # Track latency per model for smart routing
        self._latency_ema: dict[str, float] = {}

    @log_call(level="debug", include_result=True)
    def route(self, analysis: AnalysisResult, explicit_model: str | None = None) -> RoutingDecision:
        """Pick the best model for a request.

        Parameters
        ----------
        analysis : AnalysisResult
            Output of the content analyzer.
        explicit_model : str | None
            If user requested a specific model ID (e.g., ``opus-4.6``), use it.
        """
        t0 = time.monotonic()

        # 1. Explicit model passthrough
        if explicit_model:
            decision = self._try_explicit(explicit_model, analysis, t0)
            if decision:
                return decision

        # 2. Find candidate models → filter by budget → sort → build decision
        candidates = _find_candidates(analysis, self.available_providers)
        affordable = self._filter_by_budget(candidates, analysis.estimated_output_tokens)
        affordable.sort(key=lambda x: (x[1], self._latency_ema.get(x[0].id, 999)))

        return self._build_decision(affordable, analysis, t0)

    def _try_explicit(
        self, model_str: str, analysis: AnalysisResult, t0: float,
    ) -> RoutingDecision | None:
        """Try to resolve an explicit model request."""
        spec = _resolve_explicit(model_str)
        if spec and spec.provider in self.available_providers:
            return RoutingDecision(
                model=spec,
                reason=f"explicit: {model_str}",
                analysis=analysis,
                routing_time_ms=(time.monotonic() - t0) * 1000,
            )
        return None

    def _filter_by_budget(
        self, candidates: list[ModelSpec], estimated_output_tokens: int,
    ) -> list[tuple[ModelSpec, float]]:
        """Filter candidates by request cost cap and budget, with fallbacks."""
        affordable = []
        for m in candidates:
            est_cost = _estimate_cost(m, estimated_output_tokens)
            if est_cost <= self.max_request_cost and self.ledger.can_spend(est_cost):
                affordable.append((m, est_cost))

        if affordable:
            return affordable

        # Budget exceeded — fall back to free models
        free_models = [m for m in candidates if m.effective_cost == 0]
        if free_models:
            return [(free_models[0], 0.0)]

        # Last resort: cheapest anyway, but warn
        m = candidates[0]
        est = _estimate_cost(m, estimated_output_tokens)
        logger.warning("budget_exceeded", model=m.id, estimated_cost=est)
        return [(m, est)]

    def _build_decision(
        self,
        affordable: list[tuple[ModelSpec, float]],
        analysis: AnalysisResult,
        t0: float,
    ) -> RoutingDecision:
        """Build final RoutingDecision from sorted affordable list."""
        primary = affordable[0]
        fallbacks = [m for m, _ in affordable[1:4]]

        decision = RoutingDecision(
            model=primary[0],
            fallbacks=fallbacks,
            estimated_cost_usd=primary[1],
            analysis=analysis,
            routing_time_ms=(time.monotonic() - t0) * 1000,
            reason=f"tier={analysis.tier.value}, cost=${primary[1]:.4f}, signals={analysis.signals[:2]}",
        )

        logger.info(
            "route_decision",
            model=decision.model.id,
            tier=analysis.tier.value,
            cost=f"${decision.estimated_cost_usd:.4f}",
            fallbacks=[f.id for f in fallbacks],
            routing_ms=f"{decision.routing_time_ms:.1f}",
        )

        return decision

    def record_latency(self, model_id: str, latency_s: float) -> None:
        """Update exponential moving average of model latency."""
        alpha = 0.3
        prev = self._latency_ema.get(model_id, latency_s)
        self._latency_ema[model_id] = alpha * latency_s + (1 - alpha) * prev

    def record_cost(self, cost_usd: float) -> None:
        self.ledger.record(cost_usd)


def _find_candidates(analysis: AnalysisResult, available_providers: set[str]) -> list[ModelSpec]:
    """Find candidate models for the analysis tier, with fallbacks."""
    candidates = models_for_tier(
        analysis.tier,
        required_caps=analysis.capabilities if analysis.capabilities else None,
        available_providers=available_providers,
    )

    if not candidates:
        for fallback_tier in [TaskTier.STANDARD, TaskTier.OPERATIONAL, TaskTier.TRIVIAL]:
            candidates = models_for_tier(
                fallback_tier, available_providers=available_providers
            )
            if candidates:
                break

    if not candidates:
        from proxym.providers import MODELS
        local = [m for m in MODELS.values() if m.provider == "ollama"]
        if local:
            return local
        raise RuntimeError("No models available — check provider API keys")

    return candidates


def _resolve_explicit(model_str: str) -> ModelSpec | None:
    """Resolve a user-provided model string to a ModelSpec.

    First resolves model aliases (e.g., "balanced" -> "claude-sonnet-4-6"),
    then looks up the model by ID or partial match.
    """
    # Resolve model aliases first
    settings = get_settings()
    resolved = settings.resolve_model_alias(model_str)

    # Try direct ID match with resolved name
    spec = get_model(resolved)
    if spec:
        return spec

    # Try partial match (e.g., "opus-4.6" → "anthropic/opus-4.6")
    from proxym.providers import MODELS
    for mid, m in MODELS.items():
        if resolved.lower() in mid.lower() or resolved.lower() in m.display_name.lower():
            return m
    return None


def _estimate_cost(model: ModelSpec, estimated_output_tokens: int) -> float:
    """Estimate USD cost for a single request."""
    # Assume input tokens ≈ 2× output tokens (typical for code tasks)
    est_input_tokens = estimated_output_tokens * 2
    input_cost = (est_input_tokens / 1_000_000) * model.input_cost
    output_cost = (estimated_output_tokens / 1_000_000) * model.output_cost
    return input_cost + output_cost
