"""Learned route statistics for dynamic bandit-style route selection (ADR-0021).

The tracker treats each viable route (ordered chain of node shards covering a
model) as a bandit arm. Observed end-to-end tokens/sec per route is kept as a
time-decayed EWMA. Selection splits traffic between:

- **exploit**: weighted-random among *proven* routes, weight ∝ tps ** alpha
  (alpha=1.0 → a 1.5x-faster route gets 1.5x the traffic);
- **scout**: with probability `explore_share`, the least-measured unproven or
  stale route is chosen so the tracker keeps learning as the network morphs.

Staleness has two mechanisms:
- continuous: sample mass decays with `stats_half_life_seconds`, so old
  observations fade;
- abrupt: every node join/leave bumps the model's *topology epoch*; stats from
  an older epoch keep their EWMA as a prior but drop back into the scout pool
  until re-measured.

Route signatures embed node ids and shard ranges, so a node re-registering
with a different shard produces a new arm automatically.
"""

from __future__ import annotations

import random
import sqlite3
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Iterable


@dataclass(frozen=True)
class RoutingConfig:
    explore_share: float = 0.3
    weight_alpha: float = 1.0
    stats_half_life_seconds: float = 600.0
    min_sample_tokens: int = 8
    # One fresh sample has mass 1.0 and decays from there; 0.5 keeps a single
    # observation "proven" for one half-life before demoting it to the scout pool.
    min_proven_weight: float = 0.5
    max_candidate_routes: int = 8
    prune_after_seconds: float = 86400.0


@dataclass
class RouteStat:
    ewma_tps: float = 0.0
    ewma_latency_ms: float = 0.0
    weight: float = 0.0  # decayed effective sample mass
    last_sample_ts: float = 0.0
    epoch: int = 0
    samples: int = 0  # lifetime raw sample count (display only)

    def decayed_weight(self, now: float, half_life: float) -> float:
        if self.weight <= 0.0:
            return 0.0
        age = max(0.0, now - self.last_sample_ts)
        return self.weight * 0.5 ** (age / half_life)


@dataclass
class RouteCandidate:
    nodes: list[Any]
    signature: str
    prior_tps: float = 0.0


def route_signature(model_key: str, nodes: Iterable[Any]) -> str:
    hops = "->".join(
        f"{getattr(n, 'node_id', '?')}[{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}]"
        for n in nodes
    )
    return f"{model_key}|{hops}"


class RouteStatsStore:
    """Thread-safe per-route decayed throughput statistics."""

    def __init__(self, config: RoutingConfig | None = None, db_path: str | None = None) -> None:
        self.config = config or RoutingConfig()
        self._lock = threading.Lock()
        self._stats: dict[str, RouteStat] = {}
        self._epochs: dict[str, int] = {}
        self._db_path = db_path
        if db_path:
            self._init_db()
            self._load_from_db()

    def epoch(self, model_key: str) -> int:
        with self._lock:
            return self._epochs.get(model_key, 0)

    def bump_epoch(self, model_keys: Iterable[str | None]) -> None:
        """Mark the topology changed for the given model keys (node join/leave)."""
        with self._lock:
            for key in model_keys:
                if key:
                    self._epochs[key] = self._epochs.get(key, 0) + 1

    def record_sample(
        self,
        model_key: str,
        signature: str,
        tokens: int,
        elapsed_seconds: float,
        now: float | None = None,
    ) -> bool:
        """Fold one completed request into the route's EWMA.

        Returns False (and records nothing) for samples below
        `min_sample_tokens` — near-empty completions come from broken routes
        and would poison the arm with meaningless throughput values.
        """
        cfg = self.config
        if tokens < cfg.min_sample_tokens or elapsed_seconds <= 0.0:
            return False
        tps = tokens / elapsed_seconds
        ts = time.time() if now is None else now
        with self._lock:
            stat = self._stats.get(signature)
            if stat is None:
                stat = RouteStat()
                self._stats[signature] = stat
            carried = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
            total = carried + 1.0
            stat.ewma_tps = (stat.ewma_tps * carried + tps) / total
            stat.ewma_latency_ms = (stat.ewma_latency_ms * carried + elapsed_seconds * 1000.0) / total
            stat.weight = total
            stat.last_sample_ts = ts
            stat.epoch = self._epochs.get(model_key, 0)
            stat.samples += 1
        return True

    def snapshot(self, signature: str, model_key: str, now: float | None = None) -> dict:
        """Point-in-time view of one route's learned state."""
        ts = time.time() if now is None else now
        cfg = self.config
        with self._lock:
            stat = self._stats.get(signature)
            current_epoch = self._epochs.get(model_key, 0)
        if stat is None:
            return {"tps": None, "weight": 0.0, "samples": 0, "status": "unsampled"}
        weight = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
        if stat.epoch != current_epoch:
            status = "stale"
        elif weight < cfg.min_proven_weight:
            status = "decayed" if stat.samples else "unsampled"
        else:
            status = "proven"
        return {
            "tps": round(stat.ewma_tps, 4) if stat.samples else None,
            "latency_ms": round(stat.ewma_latency_ms, 3) if stat.samples else None,
            "weight": round(weight, 4),
            "samples": stat.samples,
            "status": status,
        }

    def model_rows(self, model_key: str, now: float | None = None) -> list[dict]:
        """All measured route samples, including pinned experiment routes."""
        prefix = f"{model_key}|"
        with self._lock:
            signatures = [signature for signature in self._stats if signature.startswith(prefix)]
        rows = [
            {
                "signature": signature,
                "hop_count": signature.count("->") + 1,
                **self.snapshot(signature, model_key, now=now),
            }
            for signature in signatures
        ]
        return sorted(rows, key=lambda row: (row["hop_count"], row["signature"]))

    def _init_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        con.execute(
            "CREATE TABLE IF NOT EXISTS route_stats "
            "(signature TEXT PRIMARY KEY, ewma_tps REAL NOT NULL, ewma_latency_ms REAL NOT NULL, "
            "weight REAL NOT NULL, last_sample_ts REAL NOT NULL, epoch INTEGER NOT NULL, samples INTEGER NOT NULL)"
        )
        con.execute("CREATE TABLE IF NOT EXISTS route_stat_epochs (model_key TEXT PRIMARY KEY, epoch INTEGER NOT NULL)")
        con.commit()
        con.close()

    def save_to_db(self) -> None:
        if not self._db_path:
            return
        with self._lock:
            rows = [
                (signature, stat.ewma_tps, stat.ewma_latency_ms, stat.weight, stat.last_sample_ts, stat.epoch, stat.samples)
                for signature, stat in self._stats.items()
            ]
            epochs = list(self._epochs.items())
        con = sqlite3.connect(self._db_path)
        con.executemany("INSERT OR REPLACE INTO route_stats VALUES (?,?,?,?,?,?,?)", rows)
        con.executemany("INSERT OR REPLACE INTO route_stat_epochs VALUES (?,?)", epochs)
        con.commit()
        con.close()

    def _load_from_db(self) -> None:
        con = sqlite3.connect(self._db_path)  # type: ignore[arg-type]
        rows = con.execute("SELECT signature, ewma_tps, ewma_latency_ms, weight, last_sample_ts, epoch, samples FROM route_stats").fetchall()
        epochs = con.execute("SELECT model_key, epoch FROM route_stat_epochs").fetchall()
        con.close()
        self._stats = {
            signature: RouteStat(float(tps), float(latency), float(weight), float(last_sample_ts), int(epoch), int(samples))
            for signature, tps, latency, weight, last_sample_ts, epoch, samples in rows
        }
        self._epochs = {str(model_key): int(epoch) for model_key, epoch in epochs}

    def prune(self, now: float | None = None) -> int:
        """Drop routes with no samples for `prune_after_seconds`."""
        ts = time.time() if now is None else now
        cutoff = ts - self.config.prune_after_seconds
        with self._lock:
            dead = [sig for sig, stat in self._stats.items() if stat.last_sample_ts < cutoff]
            for sig in dead:
                del self._stats[sig]
        return len(dead)


def choose_route(
    candidates: list[RouteCandidate],
    store: RouteStatsStore,
    model_key: str,
    rng: random.Random | None = None,
    now: float | None = None,
) -> tuple[RouteCandidate | None, dict]:
    """Pick a route: ε-scout among unproven arms, else weighted ∝ tps**alpha.

    Returns (candidate, decision) where decision explains the pick for logs
    and diagnostics: {"mode": "scout"|"exploit"|"prior", ...}.
    """
    if not candidates:
        return None, {"mode": "none"}
    rng = rng or random
    cfg = store.config
    proven: list[tuple[RouteCandidate, float]] = []
    scouts: list[tuple[RouteCandidate, float]] = []
    for cand in candidates:
        snap = store.snapshot(cand.signature, model_key, now=now)
        if snap["status"] == "proven":
            proven.append((cand, max(float(snap["tps"] or 0.0), 1e-6)))
        else:
            scouts.append((cand, float(snap["weight"])))
    if scouts and (not proven or rng.random() < cfg.explore_share):
        # Least-measured first so new/stale arms accumulate samples fastest;
        # tiebreak on prior estimate so plausible routes get scouted first.
        scouts.sort(key=lambda item: (item[1], -item[0].prior_tps))
        pick = scouts[0][0]
        return pick, {"mode": "scout", "signature": pick.signature}
    if proven:
        weights = [tps ** cfg.weight_alpha for _, tps in proven]
        pick = rng.choices([cand for cand, _ in proven], weights=weights, k=1)[0]
        return pick, {
            "mode": "exploit",
            "signature": pick.signature,
            "candidates": len(proven),
        }
    # No stats anywhere yet — fall back to the prior (benchmark-derived) estimate.
    weights = [max(cand.prior_tps, 1e-6) ** cfg.weight_alpha for cand in candidates]
    pick = rng.choices(candidates, weights=weights, k=1)[0]
    return pick, {"mode": "prior", "signature": pick.signature}


def route_table(
    candidates: list[RouteCandidate],
    store: RouteStatsStore,
    model_key: str,
    now: float | None = None,
) -> list[dict]:
    """Diagnostics rows: learned tps, coefficient vs best, expected traffic share."""
    cfg = store.config
    rows = []
    for cand in candidates:
        snap = store.snapshot(cand.signature, model_key, now=now)
        rows.append({"candidate": cand, **snap})
    proven = [r for r in rows if r["status"] == "proven"]
    scouts = [r for r in rows if r["status"] != "proven"]
    best_tps = max((float(r["tps"]) for r in proven), default=0.0)
    exploit_budget = 1.0 - (cfg.explore_share if scouts and proven else 0.0)
    if not proven:
        exploit_budget = 0.0
    weight_sum = sum(float(r["tps"]) ** cfg.weight_alpha for r in proven) or 1.0
    out = []
    for r in rows:
        cand: RouteCandidate = r["candidate"]
        if r["status"] == "proven":
            share = exploit_budget * (float(r["tps"]) ** cfg.weight_alpha) / weight_sum
            coefficient = round(float(r["tps"]) / best_tps, 3) if best_tps else None
        else:
            share = (
                (cfg.explore_share if proven else 1.0) / len(scouts)
                if scouts
                else 0.0
            )
            coefficient = None
        out.append({
            "signature": cand.signature,
            "hops": [
                {
                    "node_id": getattr(n, "node_id", "?"),
                    "shard": f"{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}",
                    "endpoint": getattr(n, "endpoint", "?"),
                }
                for n in cand.nodes
            ],
            "tps": r["tps"],
            "latency_ms": r["latency_ms"],
            "hop_count": len(cand.nodes),
            "coefficient": coefficient,
            "expected_share": round(share, 4),
            "samples": r["samples"],
            "weight": r["weight"],
            "status": r["status"],
            "prior_tps": round(cand.prior_tps, 4),
        })
    out.sort(key=lambda r: (-(r["tps"] or 0.0), -r["prior_tps"]))
    return out