src/database.py

"""SQLite database for generation history."""

from __future__ import annotations

import json
import sqlite3
import logging
from datetime import datetime, timedelta, timezone
from pathlib import Path

logger = logging.getLogger(__name__)

_SCHEMA = """
CREATE TABLE IF NOT EXISTS generation_history (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    job_id TEXT NOT NULL UNIQUE,
    session_id TEXT NOT NULL,
    original_filename TEXT NOT NULL,
    original_path TEXT NOT NULL,
    output_path TEXT NOT NULL,
    difficulty TEXT NOT NULL,
    num_differences INTEGER NOT NULL,
    processing_time REAL NOT NULL,
    metadata TEXT NOT NULL,
    created_at TEXT NOT NULL,
    expires_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_job_id ON generation_history(job_id);
CREATE INDEX IF NOT EXISTS idx_expires_at ON generation_history(expires_at);

CREATE TABLE IF NOT EXISTS job_status (
    job_id TEXT PRIMARY KEY,
    status TEXT NOT NULL,
    progress INTEGER NOT NULL DEFAULT 0,
    current_step TEXT NOT NULL DEFAULT '',
    error TEXT,
    result_path TEXT,
    created_at TEXT NOT NULL,
    updated_at TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_job_status_updated ON job_status(updated_at);
"""


def init_db(db_path: str) -> None:
    """Create the database and tables if they don't exist.

    Raises OSError if the database directory cannot be created.
    """
    db_file = Path(db_path)

    try:
        # Create parent directory if needed
        db_file.parent.mkdir(parents=True, exist_ok=True)
        logger.info(f"Database directory ensured: {db_file.parent}")
    except OSError as e:
        logger.error(f"Failed to create database directory {db_file.parent}: {e}")
        raise OSError(
            f"Cannot create database directory {db_file.parent}. "
            f"Please ensure FLASK_ENV=production is set for deployment."
        ) from e

    try:
        # Create database and schema
        with sqlite3.connect(db_path) as conn:
            conn.executescript(_SCHEMA)
        logger.info(f"Database initialized: {db_path}")
    except sqlite3.Error as e:
        logger.error(f"Failed to initialize database: {e}")
        raise


def save_generation(
    db_path: str,
    job_id: str,
    session_id: str,
    original_filename: str,
    original_path: str,
    output_path: str,
    difficulty: str,
    num_differences: int,
    processing_time: float,
    metadata: dict,
    expiry_hours: int = 24,
) -> int:
    """Save a generation record. Returns the row id."""
    now = datetime.now(timezone.utc)
    expires = now + timedelta(hours=expiry_hours)

    with sqlite3.connect(db_path) as conn:
        cursor = conn.execute(
            """INSERT INTO generation_history
               (job_id, session_id, original_filename, original_path,
                output_path, difficulty, num_differences, processing_time,
                metadata, created_at, expires_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            (
                job_id,
                session_id,
                original_filename,
                original_path,
                output_path,
                difficulty,
                num_differences,
                processing_time,
                json.dumps(metadata, ensure_ascii=False),
                now.isoformat(),
                expires.isoformat(),
            ),
        )
        return cursor.lastrowid


def get_generation(db_path: str, job_id: str) -> dict | None:
    """Retrieve a generation record by job_id."""
    with sqlite3.connect(db_path) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            "SELECT * FROM generation_history WHERE job_id = ?", (job_id,)
        ).fetchone()
        if row is None:
            return None
        d = dict(row)
        d["metadata"] = json.loads(d["metadata"])
        return d


def cleanup_expired(db_path: str) -> int:
    """Delete expired records. Returns count of deleted rows."""
    now = datetime.now(timezone.utc).isoformat()
    with sqlite3.connect(db_path) as conn:
        cursor = conn.execute(
            "DELETE FROM generation_history WHERE expires_at < ?", (now,)
        )
        return cursor.rowcount


# Job status persistence functions for JobManager


def save_job_status(
    db_path: str,
    job_id: str,
    status: str,
    progress: int = 0,
    current_step: str = "",
    error: str | None = None,
    result_path: str | None = None,
) -> None:
    """Save or update job status in database."""
    now = datetime.now(timezone.utc).isoformat()

    with sqlite3.connect(db_path) as conn:
        # Use INSERT OR REPLACE to handle both new and existing jobs
        conn.execute(
            """INSERT OR REPLACE INTO job_status
               (job_id, status, progress, current_step, error, result_path,
                created_at, updated_at)
               VALUES (
                   ?,
                   ?,
                   ?,
                   ?,
                   ?,
                   ?,
                   COALESCE((SELECT created_at FROM job_status WHERE job_id = ?), ?),
                   ?
               )""",
            (
                job_id,
                status,
                progress,
                current_step,
                error,
                result_path,
                job_id,
                now,
                now,
            ),
        )


def get_job_status(db_path: str, job_id: str) -> dict | None:
    """Retrieve job status by job_id."""
    with sqlite3.connect(db_path) as conn:
        conn.row_factory = sqlite3.Row
        row = conn.execute(
            "SELECT * FROM job_status WHERE job_id = ?", (job_id,)
        ).fetchone()
        if row is None:
            return None
        return dict(row)


def cleanup_old_job_status(db_path: str, hours: int = 24) -> int:
    """Delete job status records older than specified hours. Returns count of deleted rows."""
    cutoff = (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat()
    with sqlite3.connect(db_path) as conn:
        cursor = conn.execute(
            "DELETE FROM job_status WHERE updated_at < ?", (cutoff,)
        )
        return cursor.rowcount