snake-python/server/database/backend/SqliteGameplayBackend.py

"""SQLite gameplay backend with the same compact layout as PostgreSQL.

Running games use ``game_turns``/``snake_states``; finished games are packed
into one zstd blob per game in ``game_replays``.
"""

import asyncio, json, sqlite3, logging, sys
from datetime import datetime, timezone
from pathlib import Path

from server.database.backend.Template import GameplayBackendTemplate
from server.database.backend.sqlite_schema import create_schema
from server.database.compact.cache import RunningGame, RunningGameCache
from server.database.compact.codec import (
  decode_body, decode_move, decode_points, encode_body, encode_move, encode_points,
)
from server.database.compact.games import GAME_COLUMNS, normalize_game
from server.database.compact.records import (
  GameBundle, SnakeIdentity, SnakeState, TurnRecord, merge_records, records_from_game_state,
)
from server.database.compact.replay_blob import REPLAY_FORMAT, pack_replay, unpack_replay

logger = logging.getLogger(__name__)
if not logger.handlers:
  _handler = logging.StreamHandler(stream=sys.stdout)
  _handler.setFormatter(logging.Formatter(fmt="%(levelname)s %(module)s: %(message)s"))
  logger.addHandler(_handler)
  logger.propagate = False

_GAME_SELECT = """
  SELECT gid, game_id, started_at, ended_at, width, height, source, map_name,
    ruleset_name, ruleset_version, game_type, your_snake_id, your_snake_name,
    your_snake_type, your_snake_version, winner_name, winner_you, final_turn,
    status, has_replay, quality_status, quality_score, quality_tier, quality_reasons_json
  FROM games
"""

def _game_from_row(row:sqlite3.Row) -> dict:
  data = dict(row)
  reasons = data.pop("quality_reasons_json", None)
  data["quality_reasons"] = json.loads(reasons) if reasons else None
  return normalize_game(data)

class SqliteGameplayBackend(GameplayBackendTemplate):
  def __init__(self, db_path:str, busy_timeout_ms:int=5000, journal_mode:str="WAL", read_only:bool=False):
    self.db_path = db_path
    self.busy_timeout_ms = max(1000, int(busy_timeout_ms))
    self._journal_mode = journal_mode
    self._read_only = read_only
    self._running = RunningGameCache()
    if not read_only:
      self._initialize_database()

  # ── connection ─────────────────────────────────────────────────────────────

  def _connect(self) -> sqlite3.Connection:
    if self._read_only:
      connection = sqlite3.connect(f"file:{self.db_path}?mode=ro", uri=True, isolation_level=None)
      connection.row_factory = sqlite3.Row
      connection.execute("PRAGMA query_only = ON")
      return connection

    connection = sqlite3.connect(
      self.db_path,
      timeout=max(1, self.busy_timeout_ms // 1000),
      isolation_level=None,
    )
    connection.row_factory = sqlite3.Row
    connection.execute("PRAGMA foreign_keys = ON")
    connection.execute(f"PRAGMA journal_mode = {self._journal_mode}")
    connection.execute("PRAGMA synchronous = NORMAL")
    connection.execute("PRAGMA temp_store = MEMORY")
    connection.execute("PRAGMA journal_size_limit = 1048576")
    connection.execute(f"PRAGMA busy_timeout = {self.busy_timeout_ms}")
    return connection

  def _initialize_database(self) -> None:
    Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
    connection = self._connect()
    try:
      if connection.execute("PRAGMA auto_vacuum").fetchone()[0] != 1:
        connection.execute("PRAGMA auto_vacuum = FULL")
        connection.execute("VACUUM")
      create_schema(connection)
      connection.execute("PRAGMA optimize")
    finally:
      connection.close()

  # ── record helpers ─────────────────────────────────────────────────────────

  def _running_game(self, connection:sqlite3.Connection, game_id:str) -> RunningGame:
    game = self._running.get(game_id)
    if game is None:
      row = connection.execute("SELECT gid FROM games WHERE game_id = ?", (game_id,)).fetchone()
      if row is None:
        raise LookupError(f"Turn for unknown game {game_id!r}")
      game = RunningGame(gid=row["gid"])
    return game

  def _upsert_identities(self, connection:sqlite3.Connection, game_id:str, identities:list[SnakeIdentity]) -> dict[str, int]:
    gsids = {}
    for identity in identities:
      row = connection.execute("""
        INSERT INTO game_snakes (game_id, snake_id, snake_name, is_you, customizations_json, gsid)
        VALUES (?, ?, ?, ?, ?, (SELECT COALESCE(MAX(gsid), 0) + 1 FROM game_snakes))
        ON CONFLICT(game_id, snake_id) DO UPDATE SET
          snake_name = excluded.snake_name,
          is_you = excluded.is_you,
          customizations_json = excluded.customizations_json
        RETURNING gsid
        """,
        (
          game_id, identity.snake_id, identity.name, 1 if identity.is_you else 0,
          self._to_json(identity.customizations or {}),
        ),
      ).fetchone()
      gsids[identity.snake_id] = row["gsid"]
    return gsids

  def _write_rows(self, connection:sqlite3.Connection, gid:int, gsids:dict[str, int], turns:list[TurnRecord], states:list[SnakeState]) -> None:
    connection.executemany("""
      INSERT INTO game_turns (gid, turn, observed_at, my_move, thinking_json, food, hazards)
      VALUES (?, ?, ?, ?, ?, ?, ?)
      ON CONFLICT(gid, turn) DO UPDATE SET
        observed_at = excluded.observed_at,
        my_move = excluded.my_move,
        thinking_json = excluded.thinking_json,
        food = excluded.food,
        hazards = excluded.hazards
      """,
      [
        (
          gid, record.turn, self._ts_to_str(record.observed_at) or self._utc_now(),
          encode_move(record.move),
          self._to_json(record.thinking) if record.thinking is not None else None,
          encode_points(record.food), encode_points(record.hazards),
        )
        for record in turns
      ],
    )
    connection.executemany("""
      INSERT INTO snake_states (gsid, turn, health, latency, body) VALUES (?, ?, ?, ?, ?)
      ON CONFLICT(gsid, turn) DO UPDATE SET
        health = excluded.health,
        latency = excluded.latency,
        body = excluded.body
      """,
      [
        (gsids[state.snake_id], state.turn, state.health, state.latency, encode_body(state.body))
        for state in states if state.snake_id in gsids
      ],
    )

  def _load_records(self, connection:sqlite3.Connection, game_id:str, gid:int) -> tuple[list[SnakeIdentity], list[TurnRecord], list[SnakeState]]:
    identities = [
      SnakeIdentity(row["snake_id"], row["snake_name"], bool(row["is_you"]), self._from_json(row["customizations_json"]) or {})
      for row in connection.execute("""
        SELECT snake_id, snake_name, is_you, customizations_json
        FROM game_snakes WHERE game_id = ?
        """, (game_id,))
    ]
    turns = [
      TurnRecord(
        turn=row["turn"], observed_at=row["observed_at"], move=decode_move(row["my_move"]),
        thinking=self._from_json(row["thinking_json"]),
        food=decode_points(row["food"]), hazards=decode_points(row["hazards"]),
      )
      for row in connection.execute("""
        SELECT turn, observed_at, my_move, thinking_json, food, hazards
        FROM game_turns WHERE gid = ? ORDER BY turn
        """, (gid,))
    ]
    states = [
      SnakeState(row["snake_id"], row["turn"], row["health"], row["latency"], decode_body(row["body"]))
      for row in connection.execute("""
        SELECT gs.snake_id, s.turn, s.health, s.latency, s.body
        FROM game_snakes AS gs
        JOIN snake_states AS s ON s.gsid = gs.gsid
        WHERE gs.game_id = ?
        """, (game_id,))
    ]
    blob = connection.execute("SELECT format, data FROM game_replays WHERE gid = ?", (gid,)).fetchone()
    if blob is not None:
      packed_turns, packed_states = unpack_replay(blob["data"], blob["format"])
      turns, states = merge_records(packed_turns, packed_states, turns, states)
    return identities, turns, states

  def _store_replay(self, connection:sqlite3.Connection, game_id:str, gid:int, turns:list[TurnRecord], states:list[SnakeState]) -> None:
    connection.execute("""
      INSERT INTO game_replays (gid, format, turn_count, data) VALUES (?, ?, ?, ?)
      ON CONFLICT(gid) DO UPDATE SET
        format = excluded.format, turn_count = excluded.turn_count, data = excluded.data
      """, (gid, REPLAY_FORMAT, len(turns), pack_replay(turns, states)))
    connection.execute("DELETE FROM game_turns WHERE gid = ?", (gid,))
    connection.execute(
      "DELETE FROM snake_states WHERE gsid IN (SELECT gsid FROM game_snakes WHERE game_id = ?)",
      (game_id,),
    )

  def _pack_game(self, connection:sqlite3.Connection, game_id:str) -> None:
    connection.execute("BEGIN IMMEDIATE")
    try:
      row = connection.execute("SELECT gid FROM games WHERE game_id = ?", (game_id,)).fetchone()
      if row is not None and connection.execute(
        "SELECT 1 FROM game_turns WHERE gid = ? LIMIT 1", (row["gid"],),
      ).fetchone():
        _, turns, states = self._load_records(connection, game_id, row["gid"])
        self._store_replay(connection, game_id, row["gid"], turns, states)
      connection.execute("COMMIT")
    except Exception:
      connection.execute("ROLLBACK")
      raise

  # ── sync write methods ─────────────────────────────────────────────────────

  def _record_game_start_sync(self, game_state:dict, snake_type:str|None=None, snake_version:str|None=None) -> None:
    game = game_state.get("game", {})
    board = game_state.get("board", {})
    you = self._extract_you(game_state)
    ruleset = game.get("ruleset", {})
    game_type = self._derive_game_type(board, ruleset)

    connection = self._connect()
    try:
      connection.execute("""
        INSERT INTO games (
          game_id, started_at, width, height, source, map_name,
          ruleset_name, ruleset_version, your_snake_id, your_snake_name,
          your_snake_type, your_snake_version, game_type, status, gid
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'running',
          (SELECT COALESCE(MAX(gid), 0) + 1 FROM games))
        ON CONFLICT(game_id) DO UPDATE SET
          width = excluded.width,
          height = excluded.height,
          source = excluded.source,
          map_name = excluded.map_name,
          ruleset_name = excluded.ruleset_name,
          ruleset_version = excluded.ruleset_version,
          your_snake_id = excluded.your_snake_id,
          your_snake_name = excluded.your_snake_name,
          your_snake_type = excluded.your_snake_type,
          your_snake_version = excluded.your_snake_version,
          game_type = excluded.game_type,
          status = 'running'
        """,
        (
          game.get("id"), self._utc_now(), board.get("width"), board.get("height"),
          game.get("source"), game.get("map"), ruleset.get("name"), ruleset.get("version"),
          you.get("id"), you.get("name"), snake_type, snake_version, game_type,
        ),
      )
      connection.execute("PRAGMA wal_checkpoint(PASSIVE)")
    finally:
      connection.close()

  def _record_turn_sync(self, game_state:dict, my_move:str|None, my_thinking:dict|None) -> None:
    game_id = game_state.get("game", {}).get("id")
    connection = self._connect()
    try:
      connection.execute("BEGIN IMMEDIATE")
      try:
        running = self._running_game(connection, game_id)
        record, states, identities, points = records_from_game_state(
          game_state, my_move, my_thinking, running.points,
        )
        if any(identity.snake_id not in running.snake_ids for identity in identities):
          running.snake_ids.update(self._upsert_identities(connection, game_id, identities))
        self._write_rows(connection, running.gid, running.snake_ids, [record], states)
        connection.execute(
          "UPDATE games SET final_turn = MAX(final_turn, ?) WHERE gid = ?",
          (record.turn, running.gid),
        )
        connection.execute("COMMIT")
      except Exception:
        connection.execute("ROLLBACK")
        raise
    finally:
      connection.close()
    running.points = points
    self._running.put(game_id, running)

  def _record_game_end_sync(self, game_state:dict) -> None:
    game_id, winner_name, winner_you, turn = self._extract_game_end_params(game_state)

    connection = self._connect()
    try:
      connection.execute("""
        UPDATE games
        SET ended_at = ?,
          winner_name = ?,
          winner_you = ?,
          final_turn = MAX(final_turn, ?),
          status = 'finished'
        WHERE game_id = ?
        """,
        (self._utc_now(), winner_name, 1 if winner_you else 0, turn, game_id),
      )
      self._pack_game(connection, game_id)
    finally:
      connection.close()
    self._running.drop(game_id)

  def _finalize_stale_running_games_sync(self, stale_after_seconds:int=600) -> int:
    threshold = max(60, int(stale_after_seconds))
    now_utc = datetime.now(timezone.utc)
    finalized = 0

    connection = self._connect()
    try:
      rows = connection.execute("""
        SELECT gid, game_id, started_at, final_turn, your_snake_id
        FROM games
        WHERE status = 'running'
        ORDER BY started_at ASC
      """).fetchall()

      for row in rows:
        started_at = self._parse_utc_timestamp(row["started_at"])
        if started_at is None or (now_utc - started_at).total_seconds() < threshold:
          continue

        game_id = row["game_id"]
        identities, _, states = self._load_records(connection, game_id, row["gid"])
        final_turn, survivors = self._survivors(identities, states, int(row["final_turn"] or 0))
        winner_you, survivor_name = self._calculate_survivor(survivors, row["your_snake_id"])

        result = connection.execute("""
          UPDATE games
          SET ended_at = ?,
            winner_name = ?,
            winner_you = ?,
            final_turn = MAX(final_turn, ?),
            status = 'finished'
          WHERE game_id = ? AND status = 'running'
          """,
          (self._utc_now(), survivor_name, 1 if winner_you else 0, final_turn, game_id),
        )
        if result.rowcount > 0:
          finalized += 1
          self._pack_game(connection, game_id)
          self._running.drop(game_id)

      late = connection.execute("""
        SELECT DISTINCT g.game_id
        FROM game_turns AS t JOIN games AS g ON g.gid = t.gid
        WHERE g.status = 'finished'
      """).fetchall()
      for late_row in late:
        self._pack_game(connection, late_row["game_id"])
    finally:
      connection.close()

    return finalized

  # ── sync read methods ──────────────────────────────────────────────────────

  def _get_summary_sync(self, recent_limit:int=15) -> dict:
    connection = self._connect()
    try:
      totals = connection.execute("""
        SELECT
          COUNT(*) AS total_games,
          SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END) AS running_games,
          SUM(CASE WHEN status = 'finished' THEN 1 ELSE 0 END) AS finished_games,
          SUM(CASE WHEN status = 'finished' AND winner_you = 1 THEN 1 ELSE 0 END) AS wins,
          SUM(CASE WHEN status = 'finished' AND winner_you = 0 THEN 1 ELSE 0 END) AS losses,
          AVG(CASE WHEN status = 'finished' THEN final_turn ELSE NULL END) AS avg_turns
        FROM games
        """).fetchone()

      by_type = connection.execute("""
        SELECT
          COALESCE(game_type, ruleset_name, 'unknown') AS type_label,
          COUNT(*) AS total,
          SUM(CASE WHEN status = 'finished' AND winner_you = 1 THEN 1 ELSE 0 END) AS wins,
          SUM(CASE WHEN status = 'finished' AND winner_you = 0 THEN 1 ELSE 0 END) AS losses
        FROM games
        WHERE status = 'finished'
        GROUP BY type_label
        ORDER BY total DESC
        """).fetchall()

      recent = connection.execute("""
        SELECT game_id, started_at, ended_at, map_name, ruleset_name, game_type,
          your_snake_name, your_snake_type, your_snake_version, winner_you, final_turn, status
        FROM games
        WHERE has_replay = 1
        ORDER BY started_at DESC
        LIMIT ?
        """,
        (max(1, int(recent_limit)),),
      ).fetchall()
    finally:
      connection.close()

    return self._build_summary_output(totals, by_type, recent)

  def _list_games_sync(self, limit:int=50) -> list[dict]:
    connection = self._connect()
    try:
      rows = connection.execute("""
        SELECT game_id, started_at, ended_at, map_name, source, ruleset_name, game_type,
          your_snake_name, your_snake_type, your_snake_version,
          winner_you, winner_name, final_turn, status
        FROM games
        WHERE has_replay = 1
        ORDER BY started_at DESC
        LIMIT ?
        """,
        (max(1, int(limit)),),
      ).fetchall()
    finally:
      connection.close()

    return self._build_game_list_output(rows)

  def _get_game_replay_sync(self, game_id:str) -> dict|None:
    connection = self._connect()
    try:
      game_row = connection.execute(_GAME_SELECT + " WHERE game_id = ? AND has_replay = 1", (game_id,)).fetchone()
      if game_row is None:
        return None
      identities, turns, states = self._load_records(connection, game_id, game_row["gid"])
    finally:
      connection.close()

    return self._build_game_replay_output(game_row, identities, turns, states)

  # ── bundles (migration / merge) ────────────────────────────────────────────

  def list_game_ids_sync(self) -> list[str]:
    connection = self._connect()
    try:
      return [row[0] for row in connection.execute("SELECT game_id FROM games ORDER BY started_at, game_id")]
    finally:
      connection.close()

  def read_game_bundle_sync(self, game_id:str) -> GameBundle|None:
    connection = self._connect()
    try:
      row = connection.execute(_GAME_SELECT + " WHERE game_id = ?", (game_id,)).fetchone()
      if row is None:
        return None
      identities, turns, states = self._load_records(connection, game_id, row["gid"])
    finally:
      connection.close()
    return GameBundle(_game_from_row(row), identities, turns, states)

  def import_game_bundle_sync(self, bundle:GameBundle) -> None:
    game = normalize_game(bundle.game)
    columns = [column for column in GAME_COLUMNS if column != "quality_reasons"] + ["quality_reasons_json"]
    values = [game[column] for column in GAME_COLUMNS if column != "quality_reasons"]
    values.append(self._to_json(game["quality_reasons"]) if game["quality_reasons"] is not None else None)
    values = [self._ts_to_str(value) if isinstance(value, datetime) else value for value in values]
    updates = ", ".join(f"{column} = excluded.{column}" for column in columns[1:])

    connection = self._connect()
    try:
      connection.execute("BEGIN IMMEDIATE")
      try:
        connection.execute(
          f"INSERT INTO games ({', '.join(columns)}, gid) "
          f"VALUES ({', '.join('?' for _ in columns)}, (SELECT COALESCE(MAX(gid), 0) + 1 FROM games)) "
          f"ON CONFLICT(game_id) DO UPDATE SET {updates}",
          values,
        )
        gid = connection.execute("SELECT gid FROM games WHERE game_id = ?", (game["game_id"],)).fetchone()["gid"]
        gsids = self._upsert_identities(connection, game["game_id"], bundle.snakes)
        if bundle.turns and game["has_replay"]:
          if game["status"] == "finished":
            self._store_replay(connection, game["game_id"], gid, bundle.turns, bundle.states)
          else:
            self._write_rows(connection, gid, gsids, bundle.turns, bundle.states)
        connection.execute("COMMIT")
      except Exception:
        connection.execute("ROLLBACK")
        raise
    finally:
      connection.close()
    self._running.drop(game["game_id"])

  # ── public async interface ─────────────────────────────────────────────────

  async def record_game_start(self, game_state:dict, snake_type:str|None=None, snake_version:str|None=None) -> None:
    await asyncio.to_thread(self._record_game_start_sync, game_state, snake_type, snake_version)

  async def record_turn(self, game_state:dict, my_move:str|None, my_thinking:dict|None=None) -> None:
    await asyncio.to_thread(self._record_turn_sync, game_state, my_move, my_thinking)

  async def record_game_end(self, game_state:dict) -> None:
    await asyncio.to_thread(self._record_game_end_sync, game_state)

  async def get_summary(self, recent_limit:int=15) -> dict:
    return await asyncio.to_thread(self._get_summary_sync, recent_limit)

  async def list_games(self, limit:int=50) -> list[dict]:
    return await asyncio.to_thread(self._list_games_sync, limit)

  async def finalize_stale_running_games(self, stale_after_seconds:int=600) -> int:
    return await asyncio.to_thread(self._finalize_stale_running_games_sync, stale_after_seconds)

  async def get_game_replay(self, game_id:str) -> dict|None:
    return await asyncio.to_thread(self._get_game_replay_sync, game_id)

  async def list_game_ids(self) -> list[str]:
    return await asyncio.to_thread(self.list_game_ids_sync)

  async def read_game_bundle(self, game_id:str) -> GameBundle|None:
    return await asyncio.to_thread(self.read_game_bundle_sync, game_id)

  async def import_game_bundle(self, bundle:GameBundle) -> None:
    await asyncio.to_thread(self.import_game_bundle_sync, bundle)

  async def close(self) -> None:
    return None