#!/usr/bin/env python3
"""Benchmark snake move latency against sampled states from gameplay SQLite."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sqlite3
from statistics import mean, median
import sys
from time import perf_counter

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from server.GameBoard import GameBoard
from snakes import SnakeBuilder

def percentile(values: list[float], quantile: float) -> float:
  ordered = sorted(values)
  index = min(len(ordered) - 1, round((len(ordered) - 1) * quantile))
  return ordered[index]

def load_states(db_path: str, samples: int, stride: int) -> list[tuple[dict, dict]]:
  connection = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
  connection.execute("PRAGMA query_only = ON")
  max_id = int(connection.execute("SELECT max(id) FROM turns").fetchone()[0] or 0)
  if max_id == 0:
    return []

  states: list[tuple[dict, dict]] = []
  next_id = max(1, max_id - (samples - 1) * stride)
  query = """
    SELECT t.id, t.board_state_json, t.you_json, t.food_json, t.hazards_json,
           g.your_snake_id, g.your_snake_name, g.width, g.height,
           g.game_id, g.source, g.map_name,
           g.ruleset_name, g.ruleset_version, t.turn
    FROM turns AS t
    JOIN games AS g ON g.game_id = t.game_id
    WHERE t.id >= ?
    ORDER BY t.id
    LIMIT 1
  """
  snake_query = """
    SELECT st.snake_id, COALESCE(gs.snake_name, st.snake_name),
           st.health, st.length, st.head_x, st.head_y, st.body_json,
           COALESCE(gs.customizations_json, '{}')
    FROM snake_turns AS st
    LEFT JOIN game_snakes AS gs
      ON gs.game_id = st.game_id AND gs.snake_id = st.snake_id
    WHERE st.game_id = ? AND st.turn = ?
    ORDER BY st.id
  """
  while len(states) < samples and next_id <= max_id:
    row = connection.execute(query, (next_id,)).fetchone()
    if row is None:
      break
    board = json.loads(row[1])
    you = json.loads(row[2])
    if not board.get("snakes"):
      snakes = []
      for snake_row in connection.execute(snake_query, (row[9], row[14])):
        snake_id = snake_row[0]
        snake_name = snake_row[1] or (row[6] if snake_id == row[5] else snake_id)
        body = json.loads(snake_row[6])
        snakes.append({
          "id": snake_id,
          "name": snake_name,
          "health": snake_row[2],
          "length": snake_row[3],
          "head": {"x": snake_row[4], "y": snake_row[5]},
          "body": body,
          "customizations": json.loads(snake_row[7]),
        })
      board = {
        "width": row[7],
        "height": row[8],
        "food": json.loads(row[3]),
        "hazards": json.loads(row[4]),
        "snakes": snakes,
      }
    if not you:
      you = next(
        (snake for snake in board.get("snakes", []) if snake.get("id") == row[5]),
        {},
      )
    if not you or not board.get("snakes"):
      next_id = int(row[0]) + stride
      continue
    metadata = {
      "game_id": row[9],
      "source": row[10] or "custom",
      "map": row[11] or "standard",
      "ruleset": {
        "name": row[12] or "standard",
        "version": row[13] or "v1.0.0",
        "settings": {},
      },
      "turn": int(row[14]),
    }
    states.append((board, {"you": you, **metadata}))
    next_id = int(row[0]) + stride
  connection.close()
  return states

def benchmark(snake_name: str, states: list[tuple[dict, dict]], repeat: int) -> dict:
  durations: list[float] = []
  moves = 0
  for pass_number in range(repeat):
    for board_data, metadata in states:
      snake = SnakeBuilder.build(snake_name)
      game_id = f"benchmark-{pass_number}-{metadata['game_id']}"
      board = GameBoard(
        game_id=game_id,
        width=board_data["width"],
        height=board_data["height"],
        ruleset=metadata["ruleset"],
        source=metadata["source"],
        map=metadata["map"],
        snake_class=snake,
      )
      state = {
        "game": {
          "id": game_id,
          "ruleset": metadata["ruleset"],
          "source": metadata["source"],
          "map": metadata["map"],
          "timeout": 500,
        },
        "turn": metadata["turn"],
        "board": board_data,
        "you": metadata["you"],
      }
      board.read_game_data(state)
      started = perf_counter()
      snake.choose_move(board)
      durations.append((perf_counter() - started) * 1000)
      moves += 1

  return {
    "snake": snake_name,
    "moves": moves,
    "mean_ms": mean(durations),
    "median_ms": median(durations),
    "p95_ms": percentile(durations, 0.95),
    "max_ms": max(durations),
  }

def main() -> None:
  parser = argparse.ArgumentParser()
  parser.add_argument("--database", required=True)
  parser.add_argument("--snake", action="append", default=[])
  parser.add_argument("--samples", type=int, default=100)
  parser.add_argument("--stride", type=int, default=997)
  parser.add_argument("--repeat", type=int, default=1)
  args = parser.parse_args()

  states = load_states(args.database, max(1, args.samples), max(1, args.stride))
  if not states:
    raise SystemExit("No gameplay states found")

  snake_names = args.snake or ["ApexBattleSnake", "PrismBattleSnake_GPT_5_6_Sol"]
  print(f"Loaded {len(states)} states from {args.database}")
  for snake_name in snake_names:
    result = benchmark(snake_name, states, max(1, args.repeat))
    print(
      f"{result['snake']}: {result['moves']} moves, "
      f"mean={result['mean_ms']:.2f} ms, median={result['median_ms']:.2f} ms, "
      f"p95={result['p95_ms']:.2f} ms, max={result['max_ms']:.2f} ms"
    )

if __name__ == "__main__":
  main()