space_app/db.py

from __future__ import annotations

import json
import sqlite3
import uuid
from collections import defaultdict
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from .models import SolveAggregate


def utcnow() -> datetime:
    return datetime.now(UTC)


def to_iso(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    return dt.isoformat()


class SessionStore:
    def __init__(self, database_path: Path) -> None:
        self.database_path = database_path
        self.database_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.database_path)
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        with self._connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS sessions (
                    id TEXT PRIMARY KEY,
                    player_name_raw TEXT NOT NULL,
                    player_name_norm TEXT NOT NULL,
                    puzzle_type TEXT NOT NULL,
                    difficulty TEXT NOT NULL,
                    dataset_variant TEXT NOT NULL,
                    puzzle_filename TEXT NOT NULL,
                    args TEXT NOT NULL,
                    status TEXT NOT NULL,
                    assigned_at TEXT NOT NULL,
                    started_at TEXT,
                    submitted_at TEXT,
                    elapsed_ms INTEGER,
                    submission_count INTEGER NOT NULL DEFAULT 0,
                    solved INTEGER NOT NULL DEFAULT 0,
                    engine TEXT NOT NULL,
                    submitted_artifact TEXT,
                    verification_payload TEXT
                )
                """
            )
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_sessions_lookup
                ON sessions (player_name_norm, puzzle_type, difficulty, solved)
                """
            )
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS submission_attempts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    attempt_n INTEGER NOT NULL,
                    submitted_at TEXT NOT NULL,
                    elapsed_ms INTEGER NOT NULL,
                    board_ascii TEXT NOT NULL,
                    solved INTEGER NOT NULL,
                    board_valid INTEGER NOT NULL,
                    verification_payload TEXT NOT NULL,
                    FOREIGN KEY (session_id) REFERENCES sessions (id)
                )
                """
            )
            conn.execute(
                """
                CREATE INDEX IF NOT EXISTS idx_attempts_session
                ON submission_attempts (session_id, attempt_n)
                """
            )

    def create_session(
        self,
        *,
        player_name_raw: str,
        player_name_norm: str,
        puzzle_type: str,
        difficulty: str,
        dataset_variant: str,
        puzzle_filename: str,
        args: str,
        engine: str,
    ) -> dict[str, Any]:
        session_id = str(uuid.uuid4())
        assigned_at = utcnow()
        record = {
            "id": session_id,
            "player_name_raw": player_name_raw,
            "player_name_norm": player_name_norm,
            "puzzle_type": puzzle_type,
            "difficulty": difficulty,
            "dataset_variant": dataset_variant,
            "puzzle_filename": puzzle_filename,
            "args": args,
            "status": "assigned",
            "assigned_at": to_iso(assigned_at),
            "started_at": None,
            "submitted_at": None,
            "elapsed_ms": None,
            "submission_count": 0,
            "solved": 0,
            "engine": engine,
            "submitted_artifact": None,
            "verification_payload": None,
        }
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO sessions (
                    id, player_name_raw, player_name_norm, puzzle_type, difficulty,
                    dataset_variant, puzzle_filename, args, status, assigned_at,
                    started_at, submitted_at, elapsed_ms, submission_count, solved,
                    engine, submitted_artifact, verification_payload
                ) VALUES (
                    :id, :player_name_raw, :player_name_norm, :puzzle_type, :difficulty,
                    :dataset_variant, :puzzle_filename, :args, :status, :assigned_at,
                    :started_at, :submitted_at, :elapsed_ms, :submission_count, :solved,
                    :engine, :submitted_artifact, :verification_payload
                )
                """,
                record,
            )
        return record

    def get_session(self, session_id: str) -> dict[str, Any] | None:
        with self._connect() as conn:
            row = conn.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
        return dict(row) if row else None

    def get_solved_filenames(
        self,
        *,
        player_name_norm: str,
        puzzle_type: str,
        difficulty: str,
    ) -> set[str]:
        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT puzzle_filename
                FROM sessions
                WHERE player_name_norm = ?
                  AND puzzle_type = ?
                  AND difficulty = ?
                  AND solved = 1
                """,
                (player_name_norm, puzzle_type, difficulty),
            ).fetchall()
        return {str(row["puzzle_filename"]) for row in rows}

    def mark_ready(self, session_id: str) -> dict[str, Any]:
        session = self.get_session(session_id)
        if session is None:
            raise KeyError(session_id)
        started_at = session["started_at"] or to_iso(utcnow())
        with self._connect() as conn:
            conn.execute(
                "UPDATE sessions SET started_at = ?, status = 'ready' WHERE id = ?",
                (started_at, session_id),
            )
        updated = self.get_session(session_id)
        assert updated is not None
        return updated

    def record_submission(
        self,
        *,
        session_id: str,
        solved: bool,
        submitted_artifact: str,
        verification_payload: dict[str, Any],
    ) -> dict[str, Any]:
        session = self.get_session(session_id)
        if session is None:
            raise KeyError(session_id)
        if not session["started_at"]:
            raise ValueError("Session has not been marked ready.")
        started_at = datetime.fromisoformat(str(session["started_at"]))
        submitted_at = utcnow()
        elapsed_ms = max(0, int((submitted_at - started_at).total_seconds() * 1000))
        submission_count = int(session["submission_count"]) + 1
        status = "solved" if solved else "attempted"
        board_valid = bool(verification_payload.get("board_valid", False))
        verification_json = json.dumps(verification_payload)
        submitted_at_iso = to_iso(submitted_at)
        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO submission_attempts (
                    session_id, attempt_n, submitted_at, elapsed_ms,
                    board_ascii, solved, board_valid, verification_payload
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    session_id,
                    submission_count,
                    submitted_at_iso,
                    elapsed_ms,
                    submitted_artifact,
                    int(solved),
                    int(board_valid),
                    verification_json,
                ),
            )
            conn.execute(
                """
                UPDATE sessions
                SET submitted_at = ?, elapsed_ms = ?, submission_count = ?, solved = ?,
                    status = ?, submitted_artifact = ?, verification_payload = ?
                WHERE id = ?
                """,
                (
                    submitted_at_iso,
                    elapsed_ms,
                    submission_count,
                    int(solved),
                    status,
                    submitted_artifact,
                    verification_json,
                    session_id,
                ),
            )
        updated = self.get_session(session_id)
        assert updated is not None
        return updated

    def list_leaderboard(
        self,
        *,
        puzzle_id: str,
        include_test: bool = False,
    ) -> list[dict[str, Any]]:
        params: list[Any] = [puzzle_id]
        where = ["puzzle_filename = ?", "started_at IS NOT NULL"]
        if not include_test:
            where.append("LOWER(player_name_norm) != 'test'")
        query = f"""
            SELECT player_name_raw, elapsed_ms, submission_count, solved, started_at, submitted_at
            FROM sessions
            WHERE {' AND '.join(where)}
            ORDER BY solved DESC,
                     CASE WHEN solved = 1 THEN elapsed_ms END ASC,
                     submission_count DESC,
                     started_at DESC
        """
        with self._connect() as conn:
            rows = conn.execute(query, params).fetchall()
        return [dict(row) for row in rows]

    def list_solves(self, filters: dict[str, str | None]) -> list[dict[str, Any]]:
        clauses = ["submitted_at IS NOT NULL"]
        params: list[str] = []
        mappings = {
            "player_name_norm": "player_name_norm = ?",
            "puzzle_type": "puzzle_type = ?",
            "difficulty": "difficulty = ?",
            "status": "status = ?",
        }
        for key, clause in mappings.items():
            value = filters.get(key)
            if value:
                clauses.append(clause)
                params.append(value)
        if filters.get("date_from"):
            clauses.append("submitted_at >= ?")
            params.append(str(filters["date_from"]))
        if filters.get("date_to"):
            clauses.append("submitted_at <= ?")
            params.append(str(filters["date_to"]))
        query = f"SELECT * FROM sessions WHERE {' AND '.join(clauses)} ORDER BY submitted_at DESC"
        with self._connect() as conn:
            rows = conn.execute(query, params).fetchall()
        results: list[dict[str, Any]] = []
        for row in rows:
            item = dict(row)
            if item.get("verification_payload"):
                item["verification_payload"] = json.loads(str(item["verification_payload"]))
            results.append(item)
        return results

    def aggregate_solves(self, filters: dict[str, str | None]) -> list[SolveAggregate]:
        solved_filters = dict(filters)
        if not solved_filters.get("status"):
            solved_filters["status"] = "solved"
        rows = self.list_solves(solved_filters)
        grouped: dict[tuple[str, str, str], list[int]] = defaultdict(list)
        for row in rows:
            elapsed_ms = row.get("elapsed_ms")
            if elapsed_ms is None:
                continue
            key = (
                str(row["player_name_norm"]),
                str(row["puzzle_type"]),
                str(row["difficulty"]),
            )
            grouped[key].append(int(elapsed_ms))
        aggregates: list[SolveAggregate] = []
        for (player_name_norm, puzzle_type, difficulty), elapsed_values in grouped.items():
            aggregates.append(
                SolveAggregate.from_rows(
                    player_name_norm=player_name_norm,
                    puzzle_type=puzzle_type,
                    difficulty=difficulty,
                    elapsed_values=elapsed_values,
                )
            )
        aggregates.sort(
            key=lambda item: (item.player_name_norm, item.puzzle_type, item.difficulty)
        )
        return aggregates