"""
Main adaptive router strategy. See README.md for design overview.

One AdaptiveRouter instance per router_name. Holds in-memory caches:
- _cells:           Beta(alpha, beta) bandit posteriors per (request_type, model)
- _session_states:  (session_key, model) -> SessionState for incremental signal updates

Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist
state and session snapshots back to Postgres.

Routing is stateless per-turn (Thompson sample fresh on every call).
"""

from __future__ import annotations

import asyncio
import time
from collections import OrderedDict
from dataclasses import asdict, dataclass
from typing import Any, Final, cast

from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
    get_last_user_message,
)
from litellm.router_strategy.adaptive_router.bandit import (
    BanditCell,
    apply_delta,
    initial_cell,
    pick_best,
)
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
from litellm.router_strategy.adaptive_router.config import (
    ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
    MIN_QUALITY_TIER_HEADER,
    MIN_QUALITY_TIER_METADATA_KEY,
    MIN_TURNS_FOR_CLEAN_CREDIT,
    OWNER_CACHE_TTL_SECONDS,
)
from litellm.router_strategy.adaptive_router.signals import (
    SessionState,
    SignalDelta,
    Turn,
    advance_session_state,
    apply_signal_delta,
    detect_response_signals,
    detect_user_feedback,
    merge_signal_deltas,
)
from litellm.router_strategy.adaptive_router.update_queue import (
    AdaptiveRouterUpdateQueue,
)
from litellm.types.utils import StandardLoggingRoutingDecision

# Sweep session-state cache when it exceeds this many live entries. Expired
# entries are dropped in bulk; amortizes to O(1) per insert.
_SESSION_STATE_SWEEP_THRESHOLD: Final[int] = 1024
_FEEDBACK_CONTEXT_MAX_ENTRIES: Final[int] = 1024
from litellm.repositories.table_repositories import AdaptiveRouterStateRepository
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import (
    AdaptiveRouterConfig,
    AdaptiveRouterPreferences,
    PreRoutingHookResponse,
    RequestType,
)


def _default_prefs() -> AdaptiveRouterPreferences:
    """Tier-2 prior with no declared strengths; used when a model omits prefs."""
    return AdaptiveRouterPreferences(quality_tier=2, strengths=[])


@dataclass(frozen=True, slots=True)
class _FeedbackContext:
    model_name: str
    request_type: RequestType
    user_content: str | None
    assistant_content: str | None
    turn_count: int
    clean_credit_awarded: bool
    expires_at: float


class AdaptiveRouter:
    """One instance per router_name. Holds in-memory caches + the update queue."""

    def __init__(
        self,
        router_name: str,
        config: AdaptiveRouterConfig,
        model_to_prefs: dict[str, AdaptiveRouterPreferences],
        model_to_cost: dict[str, float],
    ) -> None:
        self.router_name = router_name
        self.config = config
        self.model_to_prefs = model_to_prefs
        self.model_to_cost = model_to_cost
        self.queue = AdaptiveRouterUpdateQueue()

        self._cells: dict[tuple[RequestType, str], BanditCell] = {}
        self._session_states: dict[tuple[str, str], SessionState] = {}
        self._feedback_contexts: OrderedDict[str, _FeedbackContext] = OrderedDict()
        self._session_states_expiry: dict[tuple[str, str], float] = {}
        self._feedback_attributed_total: int = 0
        self._feedback_without_context_total: int = 0
        self._cross_model_feedback_total: int = 0
        self._response_signal_updates_total: int = 0
        # Set to True once the proxy flusher has loaded persisted priors from
        # Postgres. Checked to support lazy-load on hot-reloaded routers.
        self._state_loaded: bool = False
        self._lock = asyncio.Lock()

        self._init_cold_start_cells()

    # ---- Cold-start ------------------------------------------------------

    def _init_cold_start_cells(self) -> None:
        """Populate _cells with cold-start priors for every (rt, model) combination."""
        for rt in RequestType:
            for model in self.config.available_models:
                prefs = self.model_to_prefs.get(model) or _default_prefs()
                self._cells[(rt, model)] = initial_cell(prefs, rt)

    async def load_state_from_db(self, prisma_client: Any) -> None:
        """Override cold-start cells with persisted state. Called once at startup."""
        if prisma_client is None:
            return
        try:
            rows: Final = await AdaptiveRouterStateRepository(prisma_client).table.find_many(
                where={"router_name": self.router_name}
            )
            loaded = 0
            for row in rows:
                try:
                    rt = RequestType(row.request_type)
                except ValueError:
                    # Unknown taxonomy entry from an older/newer version. Skip.
                    continue
                if row.model_name not in self.config.available_models:
                    continue
                self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta)
                loaded += 1
            verbose_router_logger.info(
                "AdaptiveRouter[%s]: loaded %d cells from DB",
                self.router_name,
                loaded,
            )
        except Exception as e:
            verbose_router_logger.exception(
                "AdaptiveRouter[%s]: failed to load state from DB: %s",
                self.router_name,
                e,
            )

    # ---- Pre-routing hook ------------------------------------------------

    async def async_pre_routing_hook(
        self,
        model: str,
        request_kwargs: dict[str, Any],
        messages: list[dict[str, Any]] | None = None,
        input: str | list | None = None,
        specific_deployment: bool | None = False,
    ) -> PreRoutingHookResponse | None:
        """
        Plugin entry point invoked by `Router.async_pre_routing_hook` when the
        inbound `model` matches this adaptive router's `router_name`.

        Classifies the last user message, picks a logical model via the bandit,
        and stashes the chosen model on `request_kwargs["metadata"]` so the
        post-call hook can surface it as a response header.

        Routing is stateless per-turn: every call Thompson-samples fresh,
        regardless of any prior pick for the same session.
        """
        user_text: Final = get_last_user_message(cast(list[AllMessageValues], messages or [])) or ""

        request_type: Final = classify_prompt(user_text)
        min_quality_tier: Final = self._extract_min_quality_tier(request_kwargs)
        chosen_model: Final = await self.pick_model(request_type=request_type, min_quality_tier=min_quality_tier)
        verbose_router_logger.debug(
            "AdaptiveRouter[%s]: classified=%s -> chose %s",
            self.router_name,
            request_type.value,
            chosen_model,
        )

        # Relay the chosen logical model to the post-call hook, which surfaces
        # it as the `x-litellm-adaptive-router-model` response header. We use
        # `metadata` (not a top-level kwarg) so the value doesn't leak into
        # `litellm.acompletion(**input_kwargs)`.
        kwargs_metadata: Final = request_kwargs.setdefault("metadata", {})
        if isinstance(kwargs_metadata, dict):
            kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = chosen_model

        return PreRoutingHookResponse(
            model=chosen_model,
            messages=messages,
            routing_decision=StandardLoggingRoutingDecision(
                router_model_name=self.router_name,
                router_type="adaptive",
                routed_model=chosen_model,
                cause="bandit",
                request_type=request_type.value,
            ),
        )

    # ---- Pick model ------------------------------------------------------

    async def pick_model(
        self,
        request_type: RequestType,
        min_quality_tier: int | None = None,
    ) -> str:
        """Thompson-sample across eligible models. Stateless per-turn."""
        eligible: Final = self._eligible_models(min_quality_tier)
        if not eligible:
            raise ValueError(f"AdaptiveRouter[{self.router_name}]: no models meet min_quality_tier={min_quality_tier}")

        cells: Final = {m: self._cells[(request_type, m)] for m in eligible}
        costs: Final = {m: self.model_to_cost.get(m, 0.0) for m in eligible}
        return pick_best(
            cells,
            costs,
            quality_weight=self.config.weights.quality,
            cost_weight=self.config.weights.cost,
        )

    async def get_state_snapshot(self) -> dict[str, Any]:
        """In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
        cells: Final = []
        for (rt, model), cell in sorted(self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])):
            total = cell.alpha + cell.beta
            cells.append(
                {
                    "request_type": rt.value,
                    "model": model,
                    "alpha": cell.alpha,
                    "beta": cell.beta,
                    # Net observations that have moved the posterior, excluding
                    # the cold-start prior mass. `alpha + beta` would show the
                    # initial COLD_START_MASS (e.g. 10) before any real traffic
                    # arrives, which confuses operators reading the endpoint.
                    "samples": cell.total_samples,
                    "quality_mean": cell.alpha / total if total > 0 else 0.0,
                }
            )
        queue: Final = await self.queue.queue_size()
        now: Final = time.time()
        feedback_contexts_live = sum(1 for context in self._feedback_contexts.values() if context.expires_at > now)
        return {
            "router_name": self.router_name,
            "available_models": list(self.config.available_models),
            "weights": {
                "quality": self.config.weights.quality,
                "cost": self.config.weights.cost,
            },
            "model_costs": dict(self.model_to_cost),
            "cells": cells,
            "feedback_contexts_live": feedback_contexts_live,
            "feedback_attributed_total": self._feedback_attributed_total,
            "feedback_without_context_total": self._feedback_without_context_total,
            "cross_model_feedback_total": self._cross_model_feedback_total,
            "response_signal_updates_total": self._response_signal_updates_total,
            "queue": queue,
        }

    @staticmethod
    def _extract_min_quality_tier(
        request_kwargs: dict[str, Any],
    ) -> int | None:
        """Pull `min_quality_tier` from request headers or metadata.

        Precedence: headers (`x-litellm-min-quality-tier`) over metadata
        (`min_quality_tier`). Headers arrive lowercased from the proxy but we
        lookup case-insensitively to be safe. Unparseable values are ignored
        (treated as "not set") rather than raising — a bad header shouldn't
        fail the request.
        """
        headers: Final = request_kwargs.get("headers") or {}
        if isinstance(headers, dict):
            for k, v in headers.items():
                if isinstance(k, str) and k.lower() == MIN_QUALITY_TIER_HEADER:
                    try:
                        return int(v)
                    except (TypeError, ValueError):
                        return None

        metadata: Final = request_kwargs.get("metadata") or {}
        if isinstance(metadata, dict):
            raw: Final = metadata.get(MIN_QUALITY_TIER_METADATA_KEY)
            if raw is not None:
                try:
                    return int(raw)
                except (TypeError, ValueError):
                    return None
        return None

    def _eligible_models(self, min_quality_tier: int | None) -> list[str]:
        if min_quality_tier is None:
            return list(self.config.available_models)
        return [
            m
            for m in self.config.available_models
            if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier >= min_quality_tier
        ]

    # ---- Session state ---------------------------------------------------

    def get_or_create_session_state(
        self,
        session_id: str,
        model_name: str,
        request_type: RequestType,
    ) -> SessionState:
        key: Final = (session_id, model_name)
        now: Final = time.time()

        # Opportunistic bulk sweep when the cache grows past the threshold.
        # Cheap relative to the alternative of a bounded LRU — conversations
        # naturally become inactive within OWNER_CACHE_TTL_SECONDS.
        if len(self._session_states) >= _SESSION_STATE_SWEEP_THRESHOLD:
            self._evict_expired_session_states(now)

        state = self._session_states.get(key)
        if state is None:
            state = SessionState(
                session_id=session_id,
                router_name=self.router_name,
                model_name=model_name,
                classified_type=request_type.value,
            )
            self._session_states[key] = state
        self._session_states_expiry[key] = now + OWNER_CACHE_TTL_SECONDS
        return state

    def _evict_expired_session_states(self, now: float) -> None:
        """Drop session states whose TTL has passed. O(n) but amortized O(1)
        per insert thanks to `_SESSION_STATE_SWEEP_THRESHOLD`."""
        expired: Final = [k for k, exp in self._session_states_expiry.items() if exp <= now]
        for k in expired:
            self._session_states.pop(k, None)
            self._session_states_expiry.pop(k, None)

    async def record_turn(
        self,
        session_id: str,
        model_name: str,
        request_type: RequestType,
        turn: Turn,
    ) -> SignalDelta:
        """Attribute feedback to the previous response and response signals to the current model."""
        async with self._lock:
            now: Final = time.time()
            while self._feedback_contexts:
                oldest_context = next(iter(self._feedback_contexts.values()))
                if oldest_context.expires_at > now:
                    break
                self._feedback_contexts.popitem(last=False)
            previous: Final = self._feedback_contexts.pop(session_id, None)

            effective_request_type: Final = (
                previous.request_type if previous is not None and request_type == RequestType.GENERAL else request_type
            )
            current_state: Final = self.get_or_create_session_state(
                session_id,
                model_name,
                effective_request_type,
            )
            feedback_delta: Final = detect_user_feedback(
                previous.user_content if previous else None,
                turn.user_content,
                turn.tool_results,
                allow_satisfaction=(
                    previous is not None
                    and not previous.clean_credit_awarded
                    and previous.turn_count + 1 >= MIN_TURNS_FOR_CLEAN_CREDIT
                ),
            )
            previous_assistant: Final = previous.assistant_content if previous else None
            response_delta: Final = detect_response_signals(
                previous_assistant,
                turn.assistant_content,
                current_state.tool_call_history,
                turn.tool_calls,
                turn.tool_results,
                turn.response_status,
            )
            states_to_persist: Final[dict[str, SessionState]] = {model_name: current_state}
            bandit_deltas: Final[dict[tuple[RequestType, str], SignalDelta]] = {}

            if previous is not None:
                feedback_state: Final = self.get_or_create_session_state(
                    session_id,
                    previous.model_name,
                    previous.request_type,
                )
                apply_signal_delta(feedback_state, feedback_delta)
                if feedback_delta.satisfaction:
                    feedback_state.clean_credit_awarded = True
                states_to_persist[previous.model_name] = feedback_state
                if feedback_delta.any_fired():
                    self._feedback_attributed_total += 1
                    if previous.model_name != model_name:
                        self._cross_model_feedback_total += 1
                bandit_deltas[(previous.request_type, previous.model_name)] = feedback_delta
            else:
                if feedback_delta.any_fired():
                    self._feedback_without_context_total += 1
                initial_failure: Final = SignalDelta(failure=feedback_delta.failure)
                apply_signal_delta(current_state, initial_failure)
                bandit_deltas[(effective_request_type, model_name)] = initial_failure

            apply_signal_delta(current_state, response_delta)
            if self._compute_bandit_delta(response_delta) != (0.0, 0.0):
                self._response_signal_updates_total += 1
            current_key: Final = (effective_request_type, model_name)
            bandit_deltas[current_key] = merge_signal_deltas(
                bandit_deltas.get(current_key, SignalDelta()),
                response_delta,
            )
            advance_session_state(current_state, turn)

            next_turn_count: Final = (previous.turn_count if previous else 0) + 1
            clean_credit_awarded = bool((previous and previous.clean_credit_awarded) or feedback_delta.satisfaction)
            if len(self._feedback_contexts) >= _FEEDBACK_CONTEXT_MAX_ENTRIES:
                self._feedback_contexts.popitem(last=False)
            self._feedback_contexts[session_id] = _FeedbackContext(
                model_name=model_name,
                request_type=effective_request_type,
                user_content=turn.user_content,
                assistant_content=turn.assistant_content,
                turn_count=next_turn_count,
                clean_credit_awarded=clean_credit_awarded,
                expires_at=now + OWNER_CACHE_TTL_SECONDS,
            )

            for state_model, state in states_to_persist.items():
                await self.queue.add_session_state(
                    session_id,
                    self.router_name,
                    state_model,
                    self._persistable_session_snapshot(state),
                )

            combined_delta = SignalDelta()
            for (attribution_type, target_model), delta in bandit_deltas.items():
                combined_delta = merge_signal_deltas(combined_delta, delta)
                d_alpha, d_beta = self._compute_bandit_delta(delta)
                if d_alpha == 0 and d_beta == 0:
                    continue
                cell_key = (attribution_type, target_model)
                self._cells[cell_key] = apply_delta(
                    self._cells[cell_key],
                    d_alpha,
                    d_beta,
                )
                await self.queue.add_state_delta(
                    self.router_name,
                    attribution_type.value,
                    target_model,
                    d_alpha,
                    d_beta,
                )

            verbose_router_logger.debug(
                "AdaptiveRouter[%s]: feedback_target=%s current_model=%s delta=%s",
                self.router_name,
                previous.model_name if previous else None,
                model_name,
                combined_delta,
            )
            return combined_delta

    @staticmethod
    def _persistable_session_snapshot(state: SessionState) -> dict[str, Any]:
        snapshot: Final = asdict(state)
        for sensitive in (
            "last_user_content",
            "last_assistant_content",
            "tool_call_history",
            "pending_tool_calls",
        ):
            snapshot.pop(sensitive, None)
        return snapshot

    @staticmethod
    def _compute_bandit_delta(delta: SignalDelta) -> tuple[float, float]:
        """
        Translate per-turn signal deltas into bandit-cell deltas.

        v0 mapping (UNVALIDATED — D6):
        - satisfaction               -> +1 alpha
        - misalignment, stagnation,
          disengagement, failure     -> +1 beta each
        - loop                       -> +0.5 beta (weak; could be model OR user)
        - exhaustion                 -> 0 (uptime issue, tracked separately later)
        """
        d_alpha: Final = float(delta.satisfaction)
        d_beta = float(delta.misalignment + delta.stagnation + delta.disengagement + delta.failure) + 0.5 * delta.loop
        return d_alpha, d_beta