database.py

"""
database.py - SQLite database manager for Howlite Image Organizer.
Handles all CRUD operations for images, tags, folders, and settings.
"""

import sqlite3
import os
from datetime import datetime
from typing import Optional

_conn: Optional[sqlite3.Connection] = None
_db_path: Optional[str] = None


def init_db(library_path: str):
    """Initialize the database at the given library path."""
    global _conn, _db_path
    _db_path = os.path.join(library_path, "library.db")
    _conn = sqlite3.connect(_db_path, check_same_thread=False)
    _conn.row_factory = sqlite3.Row
    _conn.execute("PRAGMA journal_mode=WAL;")
    _conn.execute("PRAGMA foreign_keys=ON;")
    _create_schema()


def _create_schema():
    """Create tables if they don't exist."""
    _conn.executescript("""
        CREATE TABLE IF NOT EXISTS images (
            id TEXT PRIMARY KEY,
            original_name TEXT NOT NULL,
            display_name TEXT NOT NULL,
            file_extension TEXT NOT NULL,
            imported_at TEXT NOT NULL,
            rating INTEGER DEFAULT 0,
            width INTEGER,
            height INTEGER,
            file_size INTEGER,
            trash INTEGER DEFAULT 0,
            ai_analyzed INTEGER DEFAULT 0
        );

        CREATE TABLE IF NOT EXISTS tags (
            name TEXT PRIMARY KEY
        );

        CREATE TABLE IF NOT EXISTS image_tags (
            image_id TEXT NOT NULL,
            tag_name TEXT NOT NULL,
            PRIMARY KEY (image_id, tag_name),
            FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE,
            FOREIGN KEY (tag_name) REFERENCES tags(name) ON DELETE CASCADE
        );

        CREATE TABLE IF NOT EXISTS folders (
            name TEXT PRIMARY KEY
        );

        CREATE TABLE IF NOT EXISTS settings (
            key TEXT PRIMARY KEY,
            value TEXT
        );

        CREATE INDEX IF NOT EXISTS idx_images_trash_imported ON images(trash, imported_at DESC);
        CREATE INDEX IF NOT EXISTS idx_images_display_name ON images(display_name);
        CREATE INDEX IF NOT EXISTS idx_image_tags_tag ON image_tags(tag_name, image_id);
    """)
    _conn.commit()
    _migrate_schema()
    _init_default_folders()


def _migrate_schema():
    """Safely add new columns to existing databases without breaking them."""
    migrations = [
        ("folders", "icon", "TEXT DEFAULT '📁'"),
        ("tags", "color", "TEXT DEFAULT ''"),
        ("tags", "bold", "INTEGER DEFAULT 0"),
        ("tags", "pinned", "INTEGER DEFAULT 0"),
    ]
    for table, column, col_def in migrations:
        try:
            _conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {col_def}")
            _conn.commit()
        except Exception:
            pass  # Column already exists


def _init_default_folders():
    cur = _conn.execute("SELECT COUNT(*) as cnt FROM folders")
    if cur.fetchone()["cnt"] == 0:
        defaults = ["Concept Art", "Illustrations", "Références 3D", "Textures & Materials", "UI Mockups"]
        for f in defaults:
            _conn.execute("INSERT OR IGNORE INTO folders (name) VALUES (?)", (f,))
        _conn.commit()


# ─── Images Operations ───────────────────────────────────────────────────────

def add_image(image_id: str, original_name: str, extension: str, width: int, height: int, file_size: int):
    """Insert a new image record into the database."""
    display_name = os.path.splitext(original_name)[0]
    _conn.execute(
        """INSERT INTO images (id, original_name, display_name, file_extension, imported_at, width, height, file_size)
           VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
        (image_id, original_name, display_name, extension,
         datetime.now().isoformat(), width, height, file_size)
    )
    _conn.commit()


def get_all_images(include_trash=False) -> list:
    """Fetch all images, optionally including trashed ones."""
    trash_filter = "" if include_trash else "WHERE trash = 0"
    cur = _conn.execute(f"SELECT * FROM images {trash_filter} ORDER BY imported_at DESC")
    return [dict(row) for row in cur.fetchall()]


def get_images_by_tag(tag_name: str) -> list:
    """Fetch all images that have a specific tag."""
    cur = _conn.execute(
        """SELECT i.* FROM images i
           JOIN image_tags it ON i.id = it.image_id
           WHERE it.tag_name = ? AND i.trash = 0
           ORDER BY i.imported_at DESC""",
        (tag_name,)
    )
    return [dict(row) for row in cur.fetchall()]


def get_images_by_tags(tag_names: list[str]) -> list:
    """Fetch all images that match ALL specified tags (AND logic)."""
    if not tag_names:
        return get_all_images()
    placeholders = ",".join("?" * len(tag_names))
    query = f"""
        SELECT i.* FROM images i
        JOIN image_tags it ON i.id = it.image_id
        WHERE it.tag_name IN ({placeholders}) AND i.trash = 0
        GROUP BY i.id
        HAVING COUNT(DISTINCT it.tag_name) = ?
        ORDER BY i.imported_at DESC
    """
    args = list(tag_names) + [len(tag_names)]
    cur = _conn.execute(query, args)
    return [dict(row) for row in cur.fetchall()]


def get_untagged_images() -> list:
    """Fetch all images that have no tags assigned."""
    cur = _conn.execute(
        """SELECT i.* FROM images i
           LEFT JOIN image_tags it ON i.id = it.image_id
           WHERE it.tag_name IS NULL AND i.trash = 0
           ORDER BY i.imported_at DESC"""
    )
    return [dict(row) for row in cur.fetchall()]


def get_trashed_images() -> list:
    """Fetch all images that are in the trash."""
    cur = _conn.execute("SELECT * FROM images WHERE trash = 1 ORDER BY imported_at DESC")
    return [dict(row) for row in cur.fetchall()]


def get_image(image_id: str) -> Optional[dict]:
    """Get a single image record by ID."""
    cur = _conn.execute("SELECT * FROM images WHERE id = ?", (image_id,))
    row = cur.fetchone()
    return dict(row) if row else None


def update_display_name(image_id: str, new_name: str):
    """Update the display name of an image."""
    _conn.execute("UPDATE images SET display_name = ? WHERE id = ?", (new_name, image_id))
    _conn.commit()


def move_to_trash(image_id: str):
    """Mark an image as trashed."""
    _conn.execute("UPDATE images SET trash = 1 WHERE id = ?", (image_id,))
    _conn.commit()


def restore_from_trash(image_id: str):
    """Restore a trashed image."""
    _conn.execute("UPDATE images SET trash = 0 WHERE id = ?", (image_id,))
    _conn.commit()


def delete_image_permanently(image_id: str):
    """Delete an image record permanently from the database."""
    _conn.execute("DELETE FROM images WHERE id = ?", (image_id,))
    _conn.commit()


def set_ai_analyzed(image_id: str, state: bool = True):
    """Set the ai_analyzed flag on an image."""
    _conn.execute("UPDATE images SET ai_analyzed = ? WHERE id = ?", (1 if state else 0, image_id))
    _conn.commit()


# ─── Tag Operations ──────────────────────────────────────────────────────────

def add_tag(tag_name: str):
    """Add a tag to the tags master list."""
    tag_name = tag_name.strip().lower()
    if tag_name:
        _conn.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (tag_name,))
        _conn.commit()


def add_tag_to_image(image_id: str, tag_name: str):
    """Assign a tag to an image, creating the tag if necessary."""
    tag_name = tag_name.strip().lower()
    if not tag_name:
        return
    add_tag(tag_name)
    _conn.execute(
        "INSERT OR IGNORE INTO image_tags (image_id, tag_name) VALUES (?, ?)",
        (image_id, tag_name)
    )
    _conn.commit()


def remove_tag_from_image(image_id: str, tag_name: str):
    """Remove a tag association from a single image."""
    _conn.execute(
        "DELETE FROM image_tags WHERE image_id = ? AND tag_name = ?",
        (image_id, tag_name)
    )
    _conn.commit()


def replace_image_tags(image_id: str, tag_names: list[str]):
    """Replace all tags for an image with a new list of tags."""
    _conn.execute("DELETE FROM image_tags WHERE image_id = ?", (image_id,))
    for tag in tag_names:
        add_tag_to_image(image_id, tag)
    _conn.commit()


def get_tags_for_image(image_id: str) -> list:
    """Get all tag names for a specific image."""
    cur = _conn.execute(
        "SELECT tag_name FROM image_tags WHERE image_id = ? ORDER BY tag_name",
        (image_id,)
    )
    return [row["tag_name"] for row in cur.fetchall()]


def get_all_tags() -> list:
    """Get all tags with their usage count across non-trashed images."""
    if _conn is None:
        return []
    cur = _conn.execute(
        """SELECT t.name, COALESCE(t.color,'') as color,
              COALESCE(t.bold,0) as bold, COALESCE(t.pinned,0) as pinned,
              COUNT(it.image_id) as count
           FROM tags t
           LEFT JOIN image_tags it ON t.name = it.tag_name
           LEFT JOIN images i ON it.image_id = i.id AND i.trash = 0
           GROUP BY t.name
           ORDER BY t.pinned DESC, count DESC, t.name ASC"""
    )
    return [dict(row) for row in cur.fetchall()]


def delete_tag_globally(tag_name: str):
    """Delete a tag globally from the database."""
    if _conn is None: return
    _conn.execute("DELETE FROM tags WHERE name = ?", (tag_name,))
    _conn.commit()


def delete_unused_tags() -> int:
    """Delete all tags that are not assigned to any non-trashed image."""
    if _conn is None: return 0
    cur = _conn.execute(
        """DELETE FROM tags WHERE name NOT IN (
               SELECT DISTINCT tag_name FROM image_tags it
               JOIN images i ON it.image_id = i.id WHERE i.trash = 0
           )"""
    )
    count = cur.rowcount
    _conn.commit()
    return count


def rename_tag_globally(old_name: str, new_name: str):
    """Rename a tag globally."""
    if _conn is None: return
    new_name = new_name.strip().lower()
    if not new_name or old_name == new_name:
        return
    add_tag(new_name)
    _conn.execute(
        "UPDATE OR IGNORE image_tags SET tag_name = ? WHERE tag_name = ?",
        (new_name, old_name)
    )
    _conn.execute("DELETE FROM tags WHERE name = ?", (old_name,))
    _conn.commit()


# ─── Folder Operations ──────────────────────────────────────────────────────

def get_all_folders() -> list[dict]:
    """Fetch all custom asset folders with their icon."""
    if _conn is None:
        return [{"name": n, "icon": "📁"} for n in
                ["Concept Art", "Illustrations", "Références 3D", "Textures & Materials", "UI Mockups"]]
    cur = _conn.execute("SELECT name, COALESCE(icon,'📁') as icon FROM folders ORDER BY name ASC")
    return [dict(row) for row in cur.fetchall()]


def add_folder(folder_name: str):
    """Add a new custom asset folder."""
    folder_name = folder_name.strip()
    if folder_name:
        _conn.execute("INSERT OR IGNORE INTO folders (name) VALUES (?)", (folder_name,))
        _conn.commit()


def delete_folder(folder_name: str):
    """Delete a custom asset folder."""
    _conn.execute("DELETE FROM folders WHERE name = ?", (folder_name,))
    _conn.commit()


def rename_folder(old_name: str, new_name: str):
    """Rename a custom asset folder."""
    new_name = new_name.strip()
    if new_name and old_name != new_name:
        _conn.execute("UPDATE OR IGNORE folders SET name = ? WHERE name = ?", (new_name, old_name))
        _conn.commit()


def set_folder_icon(folder_name: str, icon: str):
    """Set the emoji icon for a custom asset folder."""
    if _conn is None:
        return
    _conn.execute("UPDATE folders SET icon = ? WHERE name = ?", (icon, folder_name))
    _conn.commit()


def set_tag_style(tag_name: str, color: str = "", bold: int = 0, pinned: int = 0):
    """Update visual style for a tag (color hex string, bold 0/1, pinned 0/1)."""
    if _conn is None:
        return
    _conn.execute(
        "UPDATE tags SET color = ?, bold = ?, pinned = ? WHERE name = ?",
        (color, bold, pinned, tag_name)
    )
    _conn.commit()


def get_tag_style(tag_name: str) -> dict:
    """Get the current style for a tag. Returns dict with color, bold, pinned."""
    if _conn is None:
        return {"color": "", "bold": 0, "pinned": 0}
    cur = _conn.execute(
        "SELECT COALESCE(color,'') as color, COALESCE(bold,0) as bold, COALESCE(pinned,0) as pinned "
        "FROM tags WHERE name = ?",
        (tag_name,)
    )
    row = cur.fetchone()
    return dict(row) if row else {"color": "", "bold": 0, "pinned": 0}


# ─── Settings Operations ────────────────────────────────────────────────────

def get_setting(key: str, default: str = "") -> str:
    """Get a configuration setting value."""
    cur = _conn.execute("SELECT value FROM settings WHERE key = ?", (key,))
    row = cur.fetchone()
    return row["value"] if row else default


def set_setting(key: str, value: str):
    """Set a configuration setting value."""
    _conn.execute(
        "INSERT INTO settings (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = ?",
        (key, value, value)
    )
    _conn.commit()


def reset_library():
    """Clear all images, tags, and image_tag mappings."""
    _conn.executescript("""
        DELETE FROM image_tags;
        DELETE FROM images;
        DELETE FROM tags;
    """)
    _conn.commit()