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

from __future__ import annotations

import argparse
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.database.benchmark_states import (
  _build_state, is_postgresql_source, load_postgresql_states,
)
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]]:
  if is_postgresql_source(db_path):
    return load_postgresql_states(db_path, samples, stride)

  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
    snake_rows = connection.execute(snake_query, (row[9], row[14])).fetchall()
    state = _build_state(row, snake_rows)
    if state is not None:
      states.append(state)
    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,
    help="SQLite path or postgresql:// DSN",
  )
  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()