database.py

"""
database.py
-----------
Lightweight SQLite persistence layer.

Tracks: sheets (uploaded sprite sheets), sprites (every detected/extracted
sprite with its full metadata + workflow status), and categories (built-in
+ user-added custom categories).

SQLite is more than sufficient here -- this is a single-Space, low-write-
concurrency workload (interactive review by essentially one operator at a
time), and it gives us real persistence across app restarts without an
external service dependency.
"""

from __future__ import annotations

import json
import sqlite3
import threading
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field

from config import DB_PATH, DEFAULT_CATEGORIES, DEFAULT_STYLES

_lock = threading.Lock()


def _connect() -> sqlite3.Connection:
    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL;")
    return conn


@contextmanager
def get_conn():
    with _lock:
        conn = _connect()
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()


def _add_missing_columns(conn: sqlite3.Connection, table: str, columns: dict[str, str]) -> None:
    """
    Additive schema migration: add each column only if the table doesn't
    already have it. A deployed Space keeps its SQLite file on a
    persistent volume across redeploys, so the schema has to evolve in
    place rather than assuming a fresh CREATE TABLE.
    """
    existing = {row["name"] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()}
    for name, definition in columns.items():
        if name not in existing:
            conn.execute(f"ALTER TABLE {table} ADD COLUMN {name} {definition}")


def init_db() -> None:
    with get_conn() as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sheets (
                id TEXT PRIMARY KEY,
                filename TEXT NOT NULL,
                width INTEGER,
                height INTEGER,
                sprite_count INTEGER DEFAULT 0,
                created_at REAL,
                zip_batch_id TEXT
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS sprites (
                id TEXT PRIMARY KEY,
                sheet_id TEXT NOT NULL,
                source_file TEXT NOT NULL,
                sprite_index INTEGER,
                x INTEGER, y INTEGER, width INTEGER, height INTEGER,
                png_path TEXT NOT NULL,
                png_path_transparent TEXT,
                sha256 TEXT,
                phash TEXT,
                quality_json TEXT,
                confidence REAL,
                confidence_reasons TEXT,
                component_count INTEGER,
                category TEXT DEFAULT 'unknown',
                style TEXT DEFAULT 'unknown',
                tags TEXT DEFAULT '[]',
                status TEXT DEFAULT 'pending',
                is_duplicate_of TEXT,
                duplicate_similarity REAL,
                created_at REAL,
                updated_at REAL
            )
            """
        )
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS categories (
                name TEXT PRIMARY KEY,
                kind TEXT NOT NULL DEFAULT 'category'
            )
            """
        )
        # Columns added after the first release. SQLite has no
        # "ADD COLUMN IF NOT EXISTS", and an existing Space already has a
        # populated sprites.sqlite3 on its persistent volume, so each one
        # is added only when missing rather than recreating the table.
        _add_missing_columns(
            conn,
            "sprites",
            {
                # Which segmentation path produced the sprite. Provenance,
                # not a guess -- see segmentation.DetectedSprite.source.
                "segmentation_source": "TEXT DEFAULT 'gap'",
                # 1 while category/style/tags are still a machine
                # suggestion nobody has reviewed.
                "auto_tagged": "INTEGER DEFAULT 0",
                "auto_tag_confidence": "REAL",
                # train / validation / test, assigned at upload time.
                "split": "TEXT",
            },
        )

        conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_status ON sprites(status)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_sheet ON sprites(sheet_id)")
        conn.execute("CREATE INDEX IF NOT EXISTS idx_sprites_sha256 ON sprites(sha256)")

        for cat in DEFAULT_CATEGORIES:
            conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, 'category')", (cat,))
        for style in DEFAULT_STYLES:
            conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, 'style')", (style,))


def new_id(prefix: str = "sprite") -> str:
    return f"{prefix}_{uuid.uuid4().hex[:12]}"


# ---------------------------------------------------------------------------
# Sheets
# ---------------------------------------------------------------------------

def insert_sheet(sheet_id: str, filename: str, width: int, height: int, zip_batch_id: str | None = None) -> None:
    with get_conn() as conn:
        conn.execute(
            "INSERT INTO sheets (id, filename, width, height, sprite_count, created_at, zip_batch_id) VALUES (?,?,?,?,?,?,?)",
            (sheet_id, filename, width, height, 0, time.time(), zip_batch_id),
        )


def update_sheet_sprite_count(sheet_id: str, count: int) -> None:
    with get_conn() as conn:
        conn.execute("UPDATE sheets SET sprite_count = ? WHERE id = ?", (count, sheet_id))


def list_sheets() -> list[sqlite3.Row]:
    with get_conn() as conn:
        return conn.execute("SELECT * FROM sheets ORDER BY created_at DESC").fetchall()


# ---------------------------------------------------------------------------
# Sprites
# ---------------------------------------------------------------------------

@dataclass
class SpriteRecord:
    id: str
    sheet_id: str
    source_file: str
    sprite_index: int
    x: int
    y: int
    width: int
    height: int
    png_path: str
    png_path_transparent: str | None
    sha256: str
    phash: str
    quality: dict
    confidence: float
    confidence_reasons: list[str]
    component_count: int
    category: str = "unknown"
    style: str = "unknown"
    tags: list[str] = field(default_factory=list)
    status: str = "pending"
    is_duplicate_of: str | None = None
    duplicate_similarity: float | None = None
    segmentation_source: str = "gap"
    auto_tagged: bool = False
    auto_tag_confidence: float | None = None


def insert_sprite(rec: SpriteRecord) -> None:
    now = time.time()
    with get_conn() as conn:
        conn.execute(
            """
            INSERT INTO sprites (
                id, sheet_id, source_file, sprite_index, x, y, width, height,
                png_path, png_path_transparent, sha256, phash, quality_json,
                confidence, confidence_reasons, component_count, category,
                style, tags, status, is_duplicate_of, duplicate_similarity,
                segmentation_source, auto_tagged, auto_tag_confidence,
                created_at, updated_at
            ) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
            """,
            (
                rec.id, rec.sheet_id, rec.source_file, rec.sprite_index,
                rec.x, rec.y, rec.width, rec.height,
                rec.png_path, rec.png_path_transparent, rec.sha256, rec.phash,
                json.dumps(rec.quality), rec.confidence, json.dumps(rec.confidence_reasons),
                rec.component_count, rec.category, rec.style, json.dumps(rec.tags),
                rec.status, rec.is_duplicate_of, rec.duplicate_similarity,
                rec.segmentation_source, 1 if rec.auto_tagged else 0, rec.auto_tag_confidence,
                now, now,
            ),
        )


def update_sprite_status(sprite_id: str, status: str) -> None:
    with get_conn() as conn:
        conn.execute("UPDATE sprites SET status = ?, updated_at = ? WHERE id = ?", (status, time.time(), sprite_id))


def update_sprite_fields(sprite_id: str, **fields) -> None:
    if not fields:
        return
    allowed = {
        "category", "style", "tags", "status", "x", "y", "width", "height",
        "png_path", "png_path_transparent", "auto_tagged", "split",
    }
    sets = []
    values = []
    for k, v in fields.items():
        if k not in allowed:
            continue
        if k == "tags" and isinstance(v, list):
            v = json.dumps(v)
        sets.append(f"{k} = ?")
        values.append(v)
    if not sets:
        return
    sets.append("updated_at = ?")
    values.append(time.time())
    values.append(sprite_id)
    with get_conn() as conn:
        conn.execute(f"UPDATE sprites SET {', '.join(sets)} WHERE id = ?", values)


def get_sprite(sprite_id: str) -> sqlite3.Row | None:
    with get_conn() as conn:
        return conn.execute("SELECT * FROM sprites WHERE id = ?", (sprite_id,)).fetchone()


def list_sprites(status: str | None = None, sheet_id: str | None = None, limit: int = 5000) -> list[sqlite3.Row]:
    query = "SELECT * FROM sprites WHERE 1=1"
    params: list = []
    if status:
        query += " AND status = ?"
        params.append(status)
    if sheet_id:
        query += " AND sheet_id = ?"
        params.append(sheet_id)
    query += " ORDER BY sheet_id, sprite_index LIMIT ?"
    params.append(limit)
    with get_conn() as conn:
        return conn.execute(query, params).fetchall()


def select_sprite_ids(
    scope: str = "sheet",
    sheet_id: str | None = None,
    status: str | None = None,
    sprite_ids: list[str] | None = None,
) -> list[str]:
    """
    Resolve a bulk-operation scope to concrete sprite ids.

    Scopes: "ids" (an explicit list), "sheet" (everything on one sheet),
    "status" (everything in one workflow state) or "all".
    """
    if scope == "ids":
        return list(sprite_ids or [])

    query = "SELECT id FROM sprites WHERE 1=1"
    params: list = []
    if scope == "sheet":
        if not sheet_id:
            return []
        query += " AND sheet_id = ?"
        params.append(sheet_id)
    if status:
        query += " AND status = ?"
        params.append(status)
    query += " ORDER BY sheet_id, sprite_index"
    with get_conn() as conn:
        return [r["id"] for r in conn.execute(query, params).fetchall()]


def bulk_update_metadata(
    sprite_ids: list[str],
    category: str | None = None,
    style: str | None = None,
    add_tags: list[str] | None = None,
    remove_tags: list[str] | None = None,
    replace_tags: list[str] | None = None,
) -> int:
    """
    Apply the same category / style / tag change to many sprites in ONE
    transaction.

    Tags support add/remove as well as replace, because the realistic bulk
    edit is additive: a sheet of nothing but tree sprites wants "tree"
    added to whatever the auto-tagger already worked out per sprite, not
    every other tag wiped. Replace is there for when the existing tags are
    genuinely wrong.

    Anything touched here stops counting as an unreviewed machine
    suggestion, since a human just decided it.
    """
    if not sprite_ids:
        return 0

    now = time.time()
    add = [t.strip() for t in (add_tags or []) if t.strip()]
    remove = {t.strip() for t in (remove_tags or []) if t.strip()}
    touched = 0

    with get_conn() as conn:
        for sprite_id in sprite_ids:
            row = conn.execute("SELECT tags FROM sprites WHERE id = ?", (sprite_id,)).fetchone()
            if row is None:
                continue

            sets, values = [], []
            if category:
                sets.append("category = ?")
                values.append(category)
            if style:
                sets.append("style = ?")
                values.append(style)

            if replace_tags is not None:
                tags = [t.strip() for t in replace_tags if t.strip()]
            elif add or remove:
                tags = json.loads(row["tags"]) if row["tags"] else []
                tags = [t for t in tags if t not in remove]
                for tag in add:
                    if tag not in tags:
                        tags.append(tag)
            else:
                tags = None

            if tags is not None:
                sets.append("tags = ?")
                values.append(json.dumps(tags))

            if not sets:
                continue

            sets.append("auto_tagged = 0")
            sets.append("updated_at = ?")
            values.append(now)
            values.append(sprite_id)
            conn.execute(f"UPDATE sprites SET {', '.join(sets)} WHERE id = ?", values)
            touched += 1

    return touched


def existing_hashes_for_dup_check(exclude_status: tuple[str, ...] = ("rejected",)) -> list[dict]:
    placeholders = ",".join("?" * len(exclude_status))
    with get_conn() as conn:
        rows = conn.execute(
            f"SELECT id, sha256, phash FROM sprites WHERE status NOT IN ({placeholders})",
            exclude_status,
        ).fetchall()
    return [{"id": r["id"], "sha256": r["sha256"], "phash": r["phash"]} for r in rows]


def counts_by_status() -> dict:
    with get_conn() as conn:
        rows = conn.execute("SELECT status, COUNT(*) as c FROM sprites GROUP BY status").fetchall()
    out = {"pending": 0, "accepted": 0, "rejected": 0, "uploaded": 0}
    for r in rows:
        out[r["status"]] = r["c"]
    return out


def count_possible_duplicates() -> int:
    with get_conn() as conn:
        row = conn.execute(
            "SELECT COUNT(*) as c FROM sprites WHERE is_duplicate_of IS NOT NULL AND status != 'rejected'"
        ).fetchone()
    return row["c"] if row else 0


def mark_duplicate(sprite_id: str, of_sprite_id: str, similarity: float) -> None:
    with get_conn() as conn:
        conn.execute(
            "UPDATE sprites SET is_duplicate_of = ?, duplicate_similarity = ?, updated_at = ? WHERE id = ?",
            (of_sprite_id, similarity, time.time(), sprite_id),
        )


def duplicate_clusters(statuses: tuple[str, ...] = ("pending", "accepted")) -> list[list[sqlite3.Row]]:
    """
    Group sprites into clusters of mutual duplicates, following the
    `is_duplicate_of` chain transitively.

    Chains matter: a tileset's repeated floor tile is flagged against
    whichever copy happened to be staged first, so tile #40 may point at
    #12 which points at #3. Treating those as three unrelated pairs would
    make "keep one of each" keep three. Only clusters with 2+ members are
    returned; a sprite with no duplicate at all is not a cluster.
    """
    placeholders = ",".join("?" * len(statuses))
    with get_conn() as conn:
        rows = conn.execute(
            f"SELECT * FROM sprites WHERE status IN ({placeholders})", statuses
        ).fetchall()

    by_id = {r["id"]: r for r in rows}
    parent: dict[str, str] = {}

    def find(i: str) -> str:
        parent.setdefault(i, i)
        while parent[i] != i:
            parent[i] = parent.setdefault(parent[i], parent[i])
            i = parent[i]
        return i

    def union(i: str, j: str) -> None:
        ri, rj = find(i), find(j)
        if ri != rj:
            parent[ri] = rj

    for r in rows:
        other = r["is_duplicate_of"]
        # A sprite may point at one that has since been rejected or
        # deleted; that link simply drops out of the clustering.
        if other and other in by_id:
            union(r["id"], other)

    grouped: dict[str, list[sqlite3.Row]] = {}
    for r in rows:
        if r["id"] in parent:
            grouped.setdefault(find(r["id"]), []).append(r)

    return [members for members in grouped.values() if len(members) > 1]


def clear_duplicate_flag(sprite_id: str) -> None:
    with get_conn() as conn:
        conn.execute(
            "UPDATE sprites SET is_duplicate_of = NULL, duplicate_similarity = NULL, updated_at = ? WHERE id = ?",
            (time.time(), sprite_id),
        )


# ---------------------------------------------------------------------------
# Categories
# ---------------------------------------------------------------------------

def list_categories(kind: str = "category") -> list[str]:
    with get_conn() as conn:
        rows = conn.execute("SELECT name FROM categories WHERE kind = ? ORDER BY name", (kind,)).fetchall()
    return [r["name"] for r in rows]


def add_custom_category(name: str, kind: str = "category") -> None:
    name = name.strip().lower().replace(" ", "_")
    if not name:
        return
    with get_conn() as conn:
        conn.execute("INSERT OR IGNORE INTO categories (name, kind) VALUES (?, ?)", (name, kind))