"""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.
    When *redis_url* is provided (or ``REDIS_URL`` is set), spend counters are
    persisted as Redis keys with automatic TTL-based expiry:

    - ``proxym:cost:day:YYYY-MM-DD``  (TTL 48 h)
    - ``proxym:cost:month:YYYY-MM``   (TTL 35 d)

    Falls back to pure in-memory tracking when Redis is unreachable.
    """

    _KEY_PREFIX = "proxym:cost"
    _DAY_TTL = 48 * 3600       # 48 hours
    _MONTH_TTL = 35 * 86400    # 35 days

    def __init__(
        self,
        daily_limit: float = 5.0,
        monthly_limit: float = 60.0,
        redis_url: str | None = None,
    ):
        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._redis = self._connect_redis(redis_url)
        self._rotate_keys()

    @staticmethod
    def _connect_redis(redis_url: str | None):
        """Try to connect to Redis; return client or *None*."""
        url = redis_url
        if not url:
            return None
        try:
            import redis as _redis
            client = _redis.Redis.from_url(url, decode_responses=True, socket_connect_timeout=2)
            client.ping()
            logger.info("cost_ledger_redis_connected", url=url)
            return client
        except Exception as exc:  # noqa: BLE001
            logger.warning("cost_ledger_redis_unavailable", url=url, err=str(exc))
            return None

    def _redis_day_key(self) -> str:
        return f"{self._KEY_PREFIX}:day:{self._day_key}"

    def _redis_month_key(self) -> str:
        return f"{self._KEY_PREFIX}:month:{self._month_key}"

    def _load_from_redis(self) -> None:
        """Hydrate in-memory counters from Redis."""
        if not self._redis:
            return
        try:
            day_val = self._redis.get(self._redis_day_key())
            month_val = self._redis.get(self._redis_month_key())
            if day_val is not None:
                self._daily_spend = float(day_val)
            if month_val is not None:
                self._monthly_spend = float(month_val)
        except Exception:  # noqa: BLE001
            pass

    def _save_to_redis(self) -> None:
        """Persist current counters to Redis with TTL."""
        if not self._redis:
            return
        try:
            pipe = self._redis.pipeline(transaction=False)
            pipe.set(self._redis_day_key(), f"{self._daily_spend:.8f}", ex=self._DAY_TTL)
            pipe.set(self._redis_month_key(), f"{self._monthly_spend:.8f}", ex=self._MONTH_TTL)
            pipe.execute()
        except Exception:  # noqa: BLE001
            pass

    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
        self._load_from_redis()

    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
        self._save_to_redis()

    @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.

    Uses the model registry for full resolution (ID, litellm_model, partial).
    Falls back to backward-compatible alias resolution.
    """
    from proxym.providers import resolve_model
    # Direct registry resolution (handles ID, litellm_model, partial match)
    spec = resolve_model(model_str)
    if spec:
        return spec
    # Backward-compatible alias → registry ID → ModelSpec
    settings = get_settings()
    alias_target = settings.model_aliases.get(model_str)
    if alias_target:
        spec = resolve_model(alias_target)
        if spec:
            return spec
    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
