#!/usr/bin/env python3
"""Benchmark snake move latency against sampled states from gameplay SQLite."""
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,
)
lazy from server.GameBoard import GameBoard
lazy 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()