"""Latency-aware routing -- per-backend latency tracking with percentiles and trends.

Maintains rolling windows of latency samples for each backend.  Used by the
Router to factor latency into routing decisions and by the TUI dashboard to
display backend performance.
"""

from __future__ import annotations

import logging
import math
from collections import deque

from pydantic import BaseModel, ConfigDict

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------


class LatencyStats(BaseModel):
    """Latency statistics for a single backend."""

    model_config = ConfigDict(frozen=False)

    backend_type: str
    sample_count: int
    avg_ms: float
    p50_ms: float
    p95_ms: float
    p99_ms: float
    min_ms: float
    max_ms: float
    trend: str  # "improving", "stable", "degrading"


# ---------------------------------------------------------------------------
# LatencyTracker
# ---------------------------------------------------------------------------


class LatencyTracker:
    """Tracks per-backend latency for routing decisions.

    Maintains rolling averages, percentiles, and trends for each backend.
    Thread-safe for single-threaded async use (no locks needed for asyncio).

    Usage::

        tracker = LatencyTracker(window_size=100)
        tracker.record("ollama", 45.2)
        tracker.record("ollama", 38.7)
        stats = tracker.get_stats("ollama")
        fastest = tracker.get_fastest(["ollama", "openai"])
    """

    def __init__(self, window_size: int = 100) -> None:
        if window_size < 1:
            raise ValueError("window_size must be >= 1")
        self._window_size = window_size
        self._samples: dict[str, deque[float]] = {}

    # -- recording ----------------------------------------------------------

    def record(self, backend_type: str, latency_ms: float) -> None:
        """Record a latency measurement for a backend.

        Parameters
        ----------
        backend_type:
            The backend identifier (e.g. ``"ollama"``, ``"openai"``).
        latency_ms:
            The observed latency in milliseconds.  Negative values are ignored.
        """
        if latency_ms < 0:
            logger.warning("Ignoring negative latency %.2fms for %s", latency_ms, backend_type)
            return

        if backend_type not in self._samples:
            self._samples[backend_type] = deque(maxlen=self._window_size)

        self._samples[backend_type].append(latency_ms)

    # -- statistics ---------------------------------------------------------

    def get_stats(self, backend_type: str) -> LatencyStats | None:
        """Get latency statistics for a backend.

        Returns None if no samples have been recorded for this backend.
        """
        samples = self._samples.get(backend_type)
        if not samples:
            return None

        sorted_samples = sorted(samples)
        count = len(sorted_samples)

        avg = sum(sorted_samples) / count
        p50 = _percentile(sorted_samples, 50)
        p95 = _percentile(sorted_samples, 95)
        p99 = _percentile(sorted_samples, 99)

        trend = self._compute_trend_for_backend(backend_type)

        return LatencyStats(
            backend_type=backend_type,
            sample_count=count,
            avg_ms=round(avg, 2),
            p50_ms=round(p50, 2),
            p95_ms=round(p95, 2),
            p99_ms=round(p99, 2),
            min_ms=round(sorted_samples[0], 2),
            max_ms=round(sorted_samples[-1], 2),
            trend=trend,
        )

    def get_all_stats(self) -> dict[str, LatencyStats]:
        """Get stats for all tracked backends."""
        result: dict[str, LatencyStats] = {}
        for backend_type in self._samples:
            stats = self.get_stats(backend_type)
            if stats is not None:
                result[backend_type] = stats
        return result

    def get_fastest(self, backends: list[str]) -> str | None:
        """Return the backend with lowest average latency among the given list.

        Only considers backends with recorded samples.  Returns None if
        none of the given backends have any data.
        """
        best_backend: str | None = None
        best_avg: float = math.inf

        for bt in backends:
            samples = self._samples.get(bt)
            if not samples:
                continue
            avg = sum(samples) / len(samples)
            if avg < best_avg:
                best_avg = avg
                best_backend = bt

        return best_backend

    def should_prefer_for_interactive(
        self,
        backend_type: str,
        threshold_ms: float = 500.0,
    ) -> bool:
        """Check if this backend is fast enough for interactive use.

        Returns True if the p95 latency is below the threshold, indicating
        the backend can reliably serve interactive requests.

        Returns False if no samples exist or p95 exceeds the threshold.
        """
        stats = self.get_stats(backend_type)
        if stats is None:
            return False
        return stats.p95_ms <= threshold_ms

    # -- trend computation --------------------------------------------------

    def _compute_trend_for_backend(self, backend_type: str) -> str:
        """Compute trend from the original sample deque (time-ordered)."""
        samples = self._samples.get(backend_type)
        if not samples or len(samples) < 10:
            return "stable"

        sample_list = list(samples)
        third = len(sample_list) // 3

        if third == 0:
            return "stable"

        old_avg = sum(sample_list[:third]) / third
        new_avg = sum(sample_list[-third:]) / third

        # Allow 15% tolerance before declaring a trend
        if old_avg == 0:
            return "stable"

        change_ratio = (new_avg - old_avg) / old_avg

        if change_ratio > 0.15:
            return "degrading"
        elif change_ratio < -0.15:
            return "improving"
        return "stable"

    # -- reset / helpers ----------------------------------------------------

    def reset(self, backend_type: str | None = None) -> None:
        """Clear recorded samples.

        Parameters
        ----------
        backend_type:
            If provided, only clear samples for this backend.
            If None, clear all samples.
        """
        if backend_type is not None:
            self._samples.pop(backend_type, None)
        else:
            self._samples.clear()

    @property
    def tracked_backends(self) -> list[str]:
        """Return list of backends with recorded samples."""
        return [bt for bt, s in self._samples.items() if s]

    def __repr__(self) -> str:
        return f"LatencyTracker(window={self._window_size}, backends={len(self._samples)})"


# ---------------------------------------------------------------------------
# Percentile helper
# ---------------------------------------------------------------------------


def _percentile(sorted_data: list[float], pct: float) -> float:
    """Compute a percentile from pre-sorted data using linear interpolation.

    Parameters
    ----------
    sorted_data:
        A non-empty list of floats sorted in ascending order.
    pct:
        Percentile to compute (0-100).

    Returns
    -------
    float:
        The interpolated percentile value.
    """
    if not sorted_data:
        return 0.0

    n = len(sorted_data)
    if n == 1:
        return sorted_data[0]

    # Use the "exclusive" method (like numpy's default interpolation="linear")
    # rank = (pct / 100) * (n - 1)
    rank = (pct / 100.0) * (n - 1)
    lower = int(rank)
    upper = lower + 1
    fraction = rank - lower

    if upper >= n:
        return sorted_data[-1]

    return sorted_data[lower] + fraction * (sorted_data[upper] - sorted_data[lower])
