#!/usr/bin/env python3
"""Benchmark snake move latency against sampled states from a gameplay database."""
import argparse
from pathlib import Path
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 load_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 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()