"""Aggregate stored move traces into per-snake latency statistics.

Each move stores ``my_thinking.move_trace`` with server-side stage timings. The
engine-measured latency for turn ``T`` arrives with the board of turn ``T + 1``,
so both are paired here to split total latency into strategy, server overhead
and network/outside-handler time.
"""
from collections import defaultdict
from dataclasses import dataclass, field

from server.database.compact.replay_blob import unpack_replay
from server.database.compact.thinking import MASK_KEY

STAGES = ("strategy_ms", "overhead_ms", "handler_ms", "network_ms", "engine_ms")

@dataclass
class TraceGroup:
  moves: int = 0
  timeouts: int = 0
  timeouts_fast_handler: int = 0
  values: dict[str, list[float]] = field(default_factory=lambda: defaultdict(list))

def _thinking_dict(thinking: object) -> dict:
  if not isinstance(thinking, dict):
    return {}
  if thinking.get(MASK_KEY) == 0 and isinstance(thinking.get("_raw"), dict):
    return thinking["_raw"]
  return thinking

def _number(value: object) -> float | None:
  return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None

def _latency(value: object) -> float | None:
  # Compact storage keeps latency as a small int; "" (no response) is negative.
  number = _number(value)
  return number if number is not None and number >= 0 else None

def add_game(groups: dict[tuple, TraceGroup], game: dict, you_id: str, turns: list, states: list) -> None:
  """Add every traced move of one game to ``groups`` keyed by snake/version/mode."""
  latency_by_turn = {s.turn: _latency(s.latency) for s in states if s.snake_id == you_id}
  for record in turns:
    thinking = _thinking_dict(record.thinking)
    trace = thinking.get("move_trace")
    if not isinstance(trace, dict):
      continue
    handler = _number(trace.get("handler_ms"))
    strategy = _number(trace.get("strategy_ms"))
    if handler is None:
      continue
    key = (game["your_snake_type"], game["your_snake_version"], thinking.get("mode") or "fallback")
    group = groups.setdefault(key, TraceGroup())
    group.moves += 1
    group.values["handler_ms"].append(handler)
    if strategy is not None:
      group.values["strategy_ms"].append(strategy)
      group.values["overhead_ms"].append(max(0.0, handler - strategy))
    engine = latency_by_turn.get(record.turn + 1)
    if engine is None:
      continue
    group.values["engine_ms"].append(engine)
    group.values["network_ms"].append(max(0.0, engine - handler))
    timeout_ms = (_number(trace.get("budget_ms")) or 450.0) + 50.0
    if engine >= timeout_ms:
      group.timeouts += 1
      if handler < timeout_ms - 100:
        group.timeouts_fast_handler += 1

def percentile(values: list[float], quantile: float) -> float | None:
  if not values:
    return None
  ordered = sorted(values)
  return ordered[min(len(ordered) - 1, int(len(ordered) * quantile))]

def _fmt(value: float | None) -> str:
  return "-" if value is None else f"{value:.0f}"

def summary_rows(groups: dict[tuple, TraceGroup]) -> tuple[list[str], list[list]]:
  headers = ["snake", "version", "mode", "moves", "timeouts", "ext.timeouts"]
  for stage in STAGES:
    headers.append(f"{stage.removesuffix('_ms')} p50/p90/p99")
  rows = []
  for (snake, version, mode), group in sorted(groups.items(), key=lambda item: tuple(str(v) for v in item[0])):
    row = [
      snake, version, mode, group.moves,
      f"{group.timeouts} ({100.0 * group.timeouts / group.moves:.1f}%)" if group.moves else "0",
      group.timeouts_fast_handler,
    ]
    for stage in STAGES:
      values = group.values.get(stage, [])
      row.append("/".join(_fmt(percentile(values, q)) for q in (0.5, 0.9, 0.99)))
    rows.append(row)
  return headers, rows

async def load_groups(connection, where: str, args: list, max_games: int) -> tuple[dict[tuple, TraceGroup], int]:
  """Read the newest finished replays matching ``where``; never writes."""
  args = [*args, max(1, max_games)]
  rows = await connection.fetch(f"""
    SELECT g.game_id, g.your_snake_id, g.your_snake_type, g.your_snake_version,
           r.format, r.data
    FROM games AS g
    JOIN game_replays AS r ON r.gid = g.gid
    {where}
    ORDER BY g.started_at DESC
    LIMIT ${len(args)}
  """, *args)
  groups: dict[tuple, TraceGroup] = {}
  for row in rows:
    turns, states = unpack_replay(row["data"], row["format"])
    add_game(groups, dict(row), row["your_snake_id"], turns, states)
  return groups, len(rows)