database.py

"""
database.py — SQLite persistence layer (v2).

Schema auto-migrates from v1 → v2 on first import so old data is preserved.

Tables
------
  scans       — one row per image/video submission
                   (includes per-face signals, image hash, model version)
  scan_cache  — SHA-256 → results mapping for instant re-upload replay

Each scan stores a JSON blob of per-face signals so the history tab can
show rich signal breakdowns for past analyses.
"""

from __future__ import annotations

import json
import logging
import os
import sqlite3
import time
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Tuple

from config import settings

logger = logging.getLogger(__name__)

DB_PATH = settings.db_path
CACHE_DB_PATH = settings.cache_db_path

# ── Schema version tracking ────────────────────────────────────────────────
SCHEMA_VERSION = 2  # bump when making backward-incompatible changes


# ======================================================================
#  Initialisation & migration
# ======================================================================

def _get_connection(db: str = DB_PATH) -> sqlite3.Connection:
    """Return a connection with WAL mode + row factory."""
    conn = sqlite3.connect(db, detect_types=sqlite3.PARSE_DECLTYPES)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA foreign_keys=ON")
    return conn


def init_db() -> None:
    """Idempotent initialisation — creates / migrates tables."""
    conn = _get_connection(DB_PATH)
    cursor = conn.cursor()

    # ── Create (or migrate) the main scans table ──────────────────────
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS scans_v2 (
            id              INTEGER PRIMARY KEY AUTOINCREMENT,
            filename        TEXT    NOT NULL,
            timestamp       TEXT    NOT NULL,
            prediction_label TEXT   NOT NULL,
            confidence      REAL    NOT NULL,
            faces_detected  INTEGER NOT NULL DEFAULT 0,
            saved_image_path TEXT,
            image_hash      TEXT,
            model_version   TEXT,
            file_size_bytes INTEGER,
            processing_time_ms REAL,
            signals_json    TEXT,
            media_type      TEXT    DEFAULT 'image',   -- 'image' | 'video'
            notes           TEXT
        )
    """)

    # Check if old v1 table exists and migrate data if needed
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='scans'")
    has_old = cursor.fetchone() is not None

    if has_old:
        # Migrate: copy old rows to v2 table (safe even if already done)
        cursor.execute("""
            INSERT OR IGNORE INTO scans_v2
                (id, filename, timestamp, prediction_label, confidence,
                 faces_detected, saved_image_path)
            SELECT id, filename, timestamp, prediction_label, confidence,
                   faces_detected, saved_image_path
            FROM scans
        """)
        logger.info("Migrated rows from old 'scans' table to scans_v2")

    # ── Cache table (separate DB for performance) ─────────────────────
    cache_conn = _get_connection(CACHE_DB_PATH)
    cache_conn.execute("""
        CREATE TABLE IF NOT EXISTS scan_cache (
            image_hash      TEXT PRIMARY KEY,
            model_version   TEXT NOT NULL,
            results_json    TEXT NOT NULL,
            created_at      TEXT NOT NULL,
            hit_count       INTEGER DEFAULT 1
        )
    """)
    cache_conn.execute("""
        CREATE INDEX IF NOT EXISTS idx_cache_created
        ON scan_cache(created_at)
    """)
    cache_conn.commit()
    cache_conn.close()

    # ── Schema version marker ─────────────────────────────────────────
    cursor.execute(f"PRAGMA user_version = {SCHEMA_VERSION}")
    conn.commit()
    conn.close()
    logger.info("Database initialised (schema v%d)", SCHEMA_VERSION)


# ======================================================================
#  Scan CRUD
# ======================================================================

def log_scan(
    filename: str,
    label: str,
    confidence: float,
    faces_detected: int,
    saved_image_path: str = "",
    image_hash: Optional[str] = None,
    model_version: Optional[str] = None,
    file_size_bytes: Optional[int] = None,
    processing_time_ms: Optional[float] = None,
    signals: Optional[List[Dict[str, Any]]] = None,
    media_type: str = "image",
    notes: str = "",
) -> int:
    """Insert a scan record and return its ID."""
    conn = _get_connection()
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    signals_json = json.dumps(signals) if signals else None

    conn.execute("""
        INSERT INTO scans_v2
            (filename, timestamp, prediction_label, confidence, faces_detected,
             saved_image_path, image_hash, model_version, file_size_bytes,
             processing_time_ms, signals_json, media_type, notes)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        filename, timestamp, label, confidence, faces_detected,
        saved_image_path, image_hash, model_version, file_size_bytes,
        processing_time_ms, signals_json, media_type, notes,
    ))
    conn.commit()
    row_id = conn.execute("SELECT last_insert_rowid()").fetchone()[0]
    conn.close()
    return row_id


def get_scan_history(
    limit: int = 100,
    offset: int = 0,
    search: Optional[str] = None,
    label_filter: Optional[str] = None,
    media_type_filter: Optional[str] = None,
    days_back: Optional[int] = None,
    sort_by: str = "id",
    sort_dir: str = "DESC",
) -> List[Tuple]:
    """
    Retrieve scan records with optional filtering & pagination.

    Parameters
    ----------
    limit, offset       : pagination
    search              : text search across filename & notes
    label_filter        : 'Real' | 'Fake' | 'No Face Detected'
    media_type_filter   : 'image' | 'video'
    days_back           : only records from the last N days
    sort_by             : column name (id, timestamp, confidence, …)
    sort_dir            : ASC | DESC
    """
    conn = _get_connection()
    where_clauses: List[str] = []
    params: List[Any] = []

    if search:
        where_clauses.append("(filename LIKE ? OR notes LIKE ?)")
        params.extend([f"%{search}%", f"%{search}%"])
    if label_filter:
        where_clauses.append("prediction_label = ?")
        params.append(label_filter)
    if media_type_filter:
        where_clauses.append("media_type = ?")
        params.append(media_type_filter)
    if days_back is not None:
        cutoff = (datetime.now() - timedelta(days=days_back)).strftime("%Y-%m-%d %H:%M:%S")
        where_clauses.append("timestamp >= ?")
        params.append(cutoff)

    # Validate sort column to prevent SQL injection
    allowed_sorts = {"id", "timestamp", "confidence", "prediction_label", "filename", "faces_detected"}
    sort_col = sort_by if sort_by in allowed_sorts else "id"
    sort_d = "DESC" if sort_dir.upper() == "DESC" else "ASC"

    where_sql = " AND ".join(where_clauses) if where_clauses else "1"
    query = f"""
        SELECT id, filename, timestamp, prediction_label, confidence,
               faces_detected, saved_image_path, image_hash, model_version,
               file_size_bytes, processing_time_ms, signals_json, media_type
        FROM scans_v2
        WHERE {where_sql}
        ORDER BY {sort_col} {sort_d}
        LIMIT ? OFFSET ?
    """
    params.extend([limit, offset])
    rows = conn.execute(query, params).fetchall()
    conn.close()
    return [tuple(r) for r in rows]


def count_scans(
    search: Optional[str] = None,
    label_filter: Optional[str] = None,
    media_type_filter: Optional[str] = None,
    days_back: Optional[int] = None,
) -> int:
    """Return count of matching records (same filters as get_scan_history)."""
    conn = _get_connection()
    where_clauses: List[str] = []
    params: List[Any] = []

    if search:
        where_clauses.append("(filename LIKE ? OR notes LIKE ?)")
        params.extend([f"%{search}%", f"%{search}%"])
    if label_filter:
        where_clauses.append("prediction_label = ?")
        params.append(label_filter)
    if media_type_filter:
        where_clauses.append("media_type = ?")
        params.append(media_type_filter)
    if days_back:
        cutoff = (datetime.now() - timedelta(days=days_back)).strftime("%Y-%m-%d %H:%M:%S")
        where_clauses.append("timestamp >= ?")
        params.append(cutoff)

    where_sql = " AND ".join(where_clauses) if where_clauses else "1"
    count = conn.execute(
        f"SELECT COUNT(*) FROM scans_v2 WHERE {where_sql}", params
    ).fetchone()[0]
    conn.close()
    return count


def delete_scan(scan_id: int) -> bool:
    """Delete a single scan record by ID.  Returns True if a row was removed."""
    conn = _get_connection()
    cursor = conn.execute("DELETE FROM scans_v2 WHERE id = ?", (scan_id,))
    deleted = cursor.rowcount > 0
    conn.commit()
    conn.close()
    if deleted:
        logger.info("Deleted scan record id=%d", scan_id)
    return deleted


def delete_scan_by_image_hash(image_hash: str) -> int:
    """Delete all scans matching an image hash (dedup cleanup)."""
    conn = _get_connection()
    cursor = conn.execute("DELETE FROM scans_v2 WHERE image_hash = ?", (image_hash,))
    count = cursor.rowcount
    conn.commit()
    conn.close()
    if count:
        logger.info("Deleted %d scan(s) with hash %s", count, image_hash[:12])
    return count


def get_scan_by_id(scan_id: int) -> Optional[Dict[str, Any]]:
    """Return a single scan as a dict, or None."""
    conn = _get_connection()
    row = conn.execute(
        "SELECT * FROM scans_v2 WHERE id = ?", (scan_id,)
    ).fetchone()
    conn.close()
    if row is None:
        return None
    return dict(row)


def get_stats() -> Dict[str, Any]:
    """Return aggregate stats for the dashboard."""
    conn = _get_connection()
    total = conn.execute("SELECT COUNT(*) FROM scans_v2").fetchone()[0]
    fake_count = conn.execute(
        "SELECT COUNT(*) FROM scans_v2 WHERE prediction_label = 'Fake'"
    ).fetchone()[0]
    real_count = conn.execute(
        "SELECT COUNT(*) FROM scans_v2 WHERE prediction_label = 'Real'"
    ).fetchone()[0]
    noface_count = conn.execute(
        "SELECT COUNT(*) FROM scans_v2 WHERE prediction_label = 'No Face Detected'"
    ).fetchone()[0]
    video_count = conn.execute(
        "SELECT COUNT(*) FROM scans_v2 WHERE media_type = 'video'"
    ).fetchone()[0]

    # Average confidence
    row = conn.execute(
        "SELECT AVG(confidence) FROM scans_v2 WHERE prediction_label IN ('Real', 'Fake')"
    ).fetchone()
    avg_conf = round(row[0], 1) if row and row[0] else 0.0

    # Today's count
    today = datetime.now().strftime("%Y-%m-%d")
    today_count = conn.execute(
        "SELECT COUNT(*) FROM scans_v2 WHERE timestamp LIKE ?", (f"{today}%",)
    ).fetchone()[0]

    conn.close()
    return {
        "total_scans": total,
        "fake": fake_count,
        "real": real_count,
        "no_face": noface_count,
        "videos": video_count,
        "avg_confidence": avg_conf,
        "today": today_count,
        "model_version": settings.model_version,
    }


# ======================================================================
#  Image-hash cache
# ======================================================================

def get_cached_result(image_hash: str) -> Optional[Dict[str, Any]]:
    """
    Return cached analysis results for an image hash, or None.

    Cache is invalidated when model_version changes.
    """
    conn = _get_connection(CACHE_DB_PATH)
    row = conn.execute(
        "SELECT * FROM scan_cache WHERE image_hash = ?",
        (image_hash,),
    ).fetchone()
    if row is None:
        conn.close()
        return None

    row = dict(row)
    # Invalidate if model changed
    if row["model_version"] != settings.model_version:
        conn.execute("DELETE FROM scan_cache WHERE image_hash = ?", (image_hash,))
        conn.commit()
        conn.close()
        return None

    # Bump hit count
    conn.execute(
        "UPDATE scan_cache SET hit_count = hit_count + 1 WHERE image_hash = ?",
        (image_hash,),
    )
    conn.commit()
    conn.close()
    return json.loads(row["results_json"])


def set_cached_result(image_hash: str, results: Dict[str, Any]) -> None:
    """Store analysis results in cache."""
    conn = _get_connection(CACHE_DB_PATH)
    conn.execute("""
        INSERT OR REPLACE INTO scan_cache (image_hash, model_version, results_json, created_at)
        VALUES (?, ?, ?, ?)
    """, (
        image_hash,
        settings.model_version,
        json.dumps(results),
        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    ))
    conn.commit()
    conn.close()


def clear_cache(older_than_days: Optional[int] = None) -> int:
    """Clear cache entries.  Returns number of deleted rows."""
    conn = _get_connection(CACHE_DB_PATH)
    if older_than_days:
        cutoff = (datetime.now() - timedelta(days=older_than_days)).strftime("%Y-%m-%d %H:%M:%S")
        cursor = conn.execute("DELETE FROM scan_cache WHERE created_at < ?", (cutoff,))
    else:
        cursor = conn.execute("DELETE FROM scan_cache")
    count = cursor.rowcount
    conn.commit()
    conn.close()
    logger.info("Cleared %d cache entries", count)
    return count


def get_cache_stats() -> Dict[str, Any]:
    """Return cache hit/miss info."""
    conn = _get_connection(CACHE_DB_PATH)
    total = conn.execute("SELECT COUNT(*) FROM scan_cache").fetchone()[0]
    total_hits = conn.execute("SELECT SUM(hit_count) FROM scan_cache").fetchone()[0] or 0
    conn.close()
    return {"cached_images": total, "total_cache_hits": total_hits}


# ======================================================================
#  Disk usage helper
# ======================================================================

def get_disk_usage() -> Dict[str, Any]:
    """Report size of scans directory and DB files."""
    scans_size = 0
    if os.path.isdir(settings.scans_dir):
        for dirpath, _, filenames in os.walk(settings.scans_dir):
            for f in filenames:
                try:
                    scans_size += os.path.getsize(os.path.join(dirpath, f))
                except OSError:
                    pass

    db_size = 0
    for p in (settings.db_path, settings.cache_db_path):
        if os.path.exists(p):
            try:
                db_size += os.path.getsize(p)
            except OSError:
                pass

    def fmt(b: int) -> str:
        if b < 1024:
            return f"{b} B"
        elif b < 1024 ** 2:
            return f"{b / 1024:.1f} KB"
        else:
            return f"{b / 1024 ** 2:.1f} MB"

    return {
        "scans_dir": fmt(scans_size),
        "database": fmt(db_size),
        "scans_dir_bytes": scans_size,
        "database_bytes": db_size,
    }


# ======================================================================
#  Old-scans cleanup (call periodically or at startup)
# ======================================================================

def cleanup_old_scans(dry_run: bool = False) -> int:
    """
    Delete scan records older than ``max_scan_age_days`` and their
    associated image files.  Returns the number of records removed.
    """
    cutoff = (datetime.now() - timedelta(days=settings.max_scan_age_days)).strftime("%Y-%m-%d %H:%M:%S")
    conn = _get_connection()

    rows = conn.execute(
        "SELECT id, saved_image_path FROM scans_v2 WHERE timestamp < ?",
        (cutoff,),
    ).fetchall()

    removed = 0
    for row in rows:
        scan_id, img_path = row["id"], row["saved_image_path"]
        if not dry_run:
            # Delete image file
            if img_path and os.path.exists(img_path):
                try:
                    os.remove(img_path)
                except OSError as e:
                    logger.warning("Could not remove %s: %s", img_path, e)
            # Delete from DB
            conn.execute("DELETE FROM scans_v2 WHERE id = ?", (scan_id,))
            removed += 1

    if not dry_run:
        conn.commit()
    conn.close()

    if removed:
        logger.info("Cleanup: removed %d old scan(s) (cutoff=%s)", removed, cutoff)
    return removed


# ── Run init on import so callers don't have to remember ───────────────────
init_db()