Blame for 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 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'
            )
            """
        )
        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


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,
                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,
                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"}
    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 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 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))