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

from quart_common.web.env import env_bool

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

from server.database.backend.Template import GameplayBackendTemplate
from server.database.normalized_turn import compact_turn_json

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

_ZSTD_EXT = Path(os.environ.get("SQLITE_ZSTD_EXT", "/usr/local/lib/libsqlite_zstd.so")).expanduser().resolve()

class SqliteGameplayBackend(GameplayBackendTemplate):
  def __init__(self, db_path:str, busy_timeout_ms:int=5000, initialize_indexes:bool=True, journal_mode:str="WAL"):
    self.db_path = db_path
    self.busy_timeout_ms = max(1000, int(busy_timeout_ms))
    self._zstd_available = False
    self._initialize_indexes = initialize_indexes
    self._journal_mode = journal_mode
    self._initialize_database()

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

  def _connect(self) -> sqlite3.Connection:
    connection = sqlite3.connect(
      self.db_path,
      timeout=max(1, self.busy_timeout_ms // 1000),
      isolation_level=None,
    )
    connection.row_factory = sqlite3.Row

    if _ZSTD_EXT.exists() and not env_bool('DISABLE_GAMEPLAY_DB_COMPRESSION', True):
      try:
        connection.enable_load_extension(True)
        connection.load_extension(str(_ZSTD_EXT))
        self._zstd_available = True
      except sqlite3.OperationalError as e:
        logger.warning(f"sqlite-zstd extension skipped: {e}")
      finally:
        connection.enable_load_extension(False)

    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 _ensure_auto_vacuum_full(self, connection:sqlite3.Connection) -> None:
    current = connection.execute("PRAGMA auto_vacuum").fetchone()[0]
    if current != 1:
      connection.execute("PRAGMA auto_vacuum = FULL")
      connection.execute("VACUUM")

  # ── schema setup ───────────────────────────────────────────────────────────

  def _initialize_database(self) -> None:
    Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
    with self._connect() as connection:
      self._ensure_auto_vacuum_full(connection)
      connection.executescript("""
        CREATE TABLE IF NOT EXISTS games (
          game_id TEXT PRIMARY KEY,
          started_at TEXT NOT NULL,
          ended_at TEXT,
          width INTEGER,
          height INTEGER,
          source TEXT,
          map_name TEXT,
          ruleset_name TEXT,
          ruleset_version TEXT,
          your_snake_id TEXT,
          your_snake_name TEXT,
          your_snake_type TEXT,
          your_snake_version TEXT,
          winner_name TEXT,
          winner_you INTEGER NOT NULL DEFAULT 0,
          final_turn INTEGER NOT NULL DEFAULT 0,
          status TEXT NOT NULL DEFAULT 'running',
          has_replay INTEGER NOT NULL DEFAULT 1,
          quality_status TEXT NOT NULL DEFAULT 'retained',
          quality_score INTEGER,
          quality_tier TEXT,
          quality_reasons_json TEXT
        );

        CREATE TABLE IF NOT EXISTS turns (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          game_id TEXT NOT NULL,
          turn INTEGER NOT NULL,
          observed_at TEXT NOT NULL,
          my_move TEXT,
          my_thinking_json TEXT,
          board_state_json TEXT NOT NULL,
          snakes_json TEXT NOT NULL,
          you_json TEXT NOT NULL,
          food_json TEXT NOT NULL,
          hazards_json TEXT NOT NULL,
          UNIQUE (game_id, turn),
          FOREIGN KEY (game_id) REFERENCES games(game_id) ON DELETE CASCADE
        );

        CREATE TABLE IF NOT EXISTS game_snakes (
          game_id TEXT NOT NULL,
          snake_id TEXT NOT NULL,
          snake_name TEXT,
          is_you INTEGER NOT NULL DEFAULT 0,
          customizations_json TEXT NOT NULL DEFAULT '{}',
          PRIMARY KEY (game_id, snake_id),
          FOREIGN KEY (game_id) REFERENCES games(game_id) ON DELETE CASCADE
        );

        CREATE TABLE IF NOT EXISTS snake_turns (
          id INTEGER PRIMARY KEY AUTOINCREMENT,
          game_id TEXT NOT NULL,
          turn INTEGER NOT NULL,
          snake_id TEXT NOT NULL,
          snake_name TEXT,
          health INTEGER,
          length INTEGER,
          head_x INTEGER,
          head_y INTEGER,
          body_json TEXT NOT NULL,
          is_you INTEGER NOT NULL DEFAULT 0,
          inferred_move TEXT,
          UNIQUE (game_id, turn, snake_id),
          FOREIGN KEY (game_id) REFERENCES games(game_id) ON DELETE CASCADE
        );
        """)
      if self._initialize_indexes:
        self._create_indexes_if_tables(connection)
      self._ensure_column_exists(connection, "turns", "my_thinking_json", "TEXT")
      self._ensure_column_exists(connection, "games", "your_snake_type", "TEXT")
      self._ensure_column_exists(connection, "games", "your_snake_version", "TEXT")
      self._ensure_column_exists(connection, "games", "game_type", "TEXT")
      self._ensure_column_exists(connection, "snake_turns", "latency", "TEXT")
      self._ensure_column_exists(connection, "game_snakes", "customizations_json", "TEXT NOT NULL DEFAULT '{}'")
      self._ensure_column_exists(connection, "games", "winner_name", "TEXT")
      self._ensure_column_exists(connection, "games", "has_replay", "INTEGER NOT NULL DEFAULT 1")
      self._ensure_column_exists(connection, "games", "quality_status", "TEXT NOT NULL DEFAULT 'retained'")
      self._ensure_column_exists(connection, "games", "quality_score", "INTEGER")
      self._ensure_column_exists(connection, "games", "quality_tier", "TEXT")
      self._ensure_column_exists(connection, "games", "quality_reasons_json", "TEXT")
      if self._zstd_available:
        self._enable_zstd_compression(connection)
      connection.execute("PRAGMA optimize")

  def _create_indexes_if_tables(self, connection:sqlite3.Connection) -> None:
    real_tables = {
      row[0] for row in connection.execute(
        "SELECT name FROM sqlite_master WHERE type='table'"
      ).fetchall()
    }
    indexes = [
      ("idx_turns_game_turn",       "turns",       "game_id, turn"),
      ("idx_games_status",          "games",       "status"),
      ("idx_snake_turns_game_turn", "snake_turns", "game_id, turn"),
    ]
    for idx_name, table, cols in indexes:
      if table in real_tables:
        connection.execute(f"CREATE INDEX IF NOT EXISTS {idx_name} ON {table}({cols})")

  def _ensure_column_exists(self, connection:sqlite3.Connection, table_name:str, column_name:str, column_type:str) -> None:
    obj = connection.execute(
      "SELECT type FROM sqlite_master WHERE name = ?", (table_name,)
    ).fetchone()

    if obj and obj["type"] == "view":
      underlying = f"_{table_name}_zstd"
      exists = connection.execute(
        "SELECT 1 FROM sqlite_master WHERE name = ? AND type = 'table'", (underlying,)
      ).fetchone()
      if not exists:
        return
      actual_table = underlying
    else:
      actual_table = table_name

    existing = connection.execute(f"PRAGMA table_info({actual_table})").fetchall()
    if any(row["name"] == column_name for row in existing):
      return
    connection.execute(f"ALTER TABLE {actual_table} ADD COLUMN {column_name} {column_type}")

  def _enable_zstd_compression(self, connection:sqlite3.Connection) -> None:
    compressed_columns = [
      ("turns",       "board_state_json"),
      ("turns",       "snakes_json"),
      ("turns",       "you_json"),
      ("turns",       "food_json"),
      ("turns",       "hazards_json"),
      ("snake_turns", "body_json"),
    ]
    for table, column in compressed_columns:
      try:
        connection.execute(
          "SELECT zstd_enable_transparent(?)",
          [json.dumps({"table": table, "column": column, "compression_level": 6, "dict_chooser": "'a'"})],
        )
      except sqlite3.OperationalError:
        pass
    connection.execute("SELECT zstd_incremental_maintenance(null, 1)")

  # ── 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)

    with self._connect() as connection:
      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
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'running')
        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)")
      connection.execute("PRAGMA optimize")

  def _record_turn_sync(self, game_state:dict, my_move:str|None, my_thinking:dict|None) -> None:
    game = game_state.get("game", {})
    board = game_state.get("board", {})
    snakes = self._extract_snakes(game_state)
    you = self._extract_you(game_state)
    game_id = game.get("id")
    turn = int(game_state.get("turn", 0))

    board_json, snakes_json, you_json, food_json, hazards_json = compact_turn_json(board)

    with self._connect() as connection:
      connection.executemany("""
        INSERT INTO game_snakes (
          game_id, snake_id, snake_name, is_you, customizations_json
        ) VALUES (?, ?, ?, ?, ?)
        ON CONFLICT(game_id, snake_id) DO UPDATE SET
          snake_name = excluded.snake_name,
          is_you = excluded.is_you,
          customizations_json = excluded.customizations_json
        """,
        [
          (
            game_id,
            snake.get("id"),
            snake.get("name"),
            1 if snake.get("id") == you.get("id") else 0,
            self._to_json(snake.get("customizations") or {}),
          )
          for snake in snakes if snake.get("id") is not None
        ],
      )
      connection.execute("""
        INSERT INTO turns (
          game_id, turn, observed_at, my_move, my_thinking_json,
          board_state_json, snakes_json, you_json, food_json, hazards_json
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(game_id, turn) DO UPDATE SET
          observed_at = excluded.observed_at,
          my_move = excluded.my_move,
          my_thinking_json = excluded.my_thinking_json,
          board_state_json = excluded.board_state_json,
          snakes_json = excluded.snakes_json,
          you_json = excluded.you_json,
          food_json = excluded.food_json,
          hazards_json = excluded.hazards_json
        """,
        (
          game_id,
          turn,
          self._utc_now(),
          my_move,
          self._to_json(my_thinking) if my_thinking is not None else None,
          self._to_json(board_json),
          self._to_json(snakes_json),
          self._to_json(you_json),
          self._to_json(food_json),
          self._to_json(hazards_json),
        ),
      )

      previous_positions: dict[str, tuple[int, int]] = {}
      if turn > 0:
        previous_rows = connection.execute("""
          SELECT snake_id, head_x, head_y
          FROM snake_turns
          WHERE game_id = ? AND turn = ?
          """,
          (game_id, turn - 1),
        ).fetchall()
        previous_positions = {
          row["snake_id"]: (int(row["head_x"]), int(row["head_y"]))
          for row in previous_rows
          if row["head_x"] is not None and row["head_y"] is not None
        }

      you_id = you.get("id")
      for params in self._build_snake_turn_params(snakes, you_id, game_id, turn, previous_positions):
        (p_game_id, p_turn, p_snake_id, p_name, p_health, p_length,
         p_head_x, p_head_y, p_body, p_is_you, p_inferred, p_latency) = params
        connection.execute("""
          INSERT INTO snake_turns (
            game_id, turn, snake_id, snake_name, health, length,
            head_x, head_y, body_json, is_you, inferred_move, latency
          ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
          ON CONFLICT(game_id, turn, snake_id) DO UPDATE SET
            snake_name = excluded.snake_name,
            health = excluded.health,
            length = excluded.length,
            head_x = excluded.head_x,
            head_y = excluded.head_y,
            body_json = excluded.body_json,
            is_you = excluded.is_you,
            inferred_move = excluded.inferred_move,
            latency = excluded.latency
          """,
          (
            p_game_id, p_turn, p_snake_id, None, p_health, p_length,
            p_head_x, p_head_y, self._to_json(p_body),
            0,
            p_inferred, p_latency,
          ),
        )

      connection.execute("""
        UPDATE games
        SET final_turn = CASE WHEN ? > final_turn THEN ? ELSE final_turn END
        WHERE game_id = ?
        """,
        (turn, turn, game_id),
      )

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

    with self._connect() as connection:
      connection.execute("""
        UPDATE games
        SET ended_at = ?,
          winner_name = ?,
          winner_you = ?,
          final_turn = CASE WHEN ? > final_turn THEN ? ELSE final_turn END,
          status = 'finished'
        WHERE game_id = ?
        """,
        (
          self._utc_now(),
          winner_name,
          1 if winner_you else 0,
          turn,
          turn,
          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

    with self._connect() as connection:
      rows = connection.execute("""
        SELECT 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:
          continue
        if (now_utc - started_at).total_seconds() < threshold:
          continue

        game_id = row["game_id"]
        your_snake_id = row["your_snake_id"]
        final_turn = int(row["final_turn"] or 0)

        snake_rows = connection.execute("""
          SELECT st.snake_id, COALESCE(gs.snake_name, st.snake_name) AS snake_name
          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 COALESCE(gs.is_you, st.is_you) DESC, snake_name ASC
          """,
          (game_id, final_turn),
        ).fetchall()

        if len(snake_rows) == 0:
          latest_row = connection.execute(
            "SELECT MAX(turn) AS latest_turn FROM snake_turns WHERE game_id = ?",
            (game_id,),
          ).fetchone()
          if latest_row is not None and latest_row["latest_turn"] is not None:
            final_turn = int(latest_row["latest_turn"])
            snake_rows = connection.execute("""
              SELECT st.snake_id, COALESCE(gs.snake_name, st.snake_name) AS snake_name
              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 COALESCE(gs.is_you, st.is_you) DESC, snake_name ASC
              """,
              (game_id, final_turn),
            ).fetchall()

        winner_you, survivor_name = self._calculate_survivor(snake_rows, your_snake_id)

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

    return finalized

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

  def _get_summary_sync(self, recent_limit:int=15) -> dict:
    with self._connect() as connection:
      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()

    return self._build_summary_output(totals, by_type, recent)

  def _list_games_sync(self, limit:int=50) -> list[dict]:
    with self._connect() as connection:
      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()

    return self._build_game_list_output(rows)

  def _get_game_replay_sync(self, game_id:str) -> dict|None:
    with self._connect() as connection:
      game_row = connection.execute("""
        SELECT 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
        FROM games
        WHERE game_id = ? AND has_replay = 1
        """,
        (game_id,),
      ).fetchone()

      if game_row is None:
        return None

      turn_rows = connection.execute("""
        SELECT turn, observed_at, my_move, my_thinking_json,
          board_state_json, food_json, hazards_json, you_json
        FROM turns
        WHERE game_id = ?
        ORDER BY turn ASC
        """,
        (game_id,),
      ).fetchall()

      snake_rows = connection.execute("""
        SELECT st.turn, st.snake_id,
          COALESCE(gs.snake_name, st.snake_name) AS snake_name,
          st.health, st.length, st.head_x, st.head_y, st.body_json,
          COALESCE(gs.is_you, st.is_you) AS is_you,
          COALESCE(gs.customizations_json, '{}') AS customizations_json,
          st.inferred_move, st.latency
        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 = ?
        ORDER BY st.turn ASC, is_you DESC, snake_name ASC
        """,
        (game_id,),
      ).fetchall()

    return self._build_game_replay_output(game_row, turn_rows, snake_rows, decode_json=self._from_json)

  # ── 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 close(self) -> None:
    return None