# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# Licensed under the MIT License

"""
Lightweight SQLite migration system.

Each database tracks its own schema version in a ``_schema_version`` table.
Migrations are plain Python callables registered per database name.

Usage::

    from familiar.core.migrations import register_migration, run_migrations

    # Register migrations in order (version must be consecutive starting at 1)
    register_migration("files", 1, _files_v1_add_mentions_table)
    register_migration("files", 2, _files_v2_add_checksum_index)

    # At init time, run all pending migrations
    run_migrations(conn, "files")

Migration functions receive a ``sqlite3.Connection`` (inside a transaction)
and must NOT commit or rollback — the runner handles that.
"""

import logging
import sqlite3
from typing import Callable, Dict, List, Tuple

logger = logging.getLogger(__name__)

# Registry: db_name -> sorted list of (version, callable)
_registry: Dict[str, List[Tuple[int, Callable[[sqlite3.Connection], None]]]] = {}


def register_migration(
    db_name: str,
    version: int,
    fn: Callable[[sqlite3.Connection], None],
) -> None:
    """Register a migration for a database.

    Args:
        db_name: Logical database name (e.g. "files", "analytics", "audit").
        version: Integer version this migration upgrades TO (must be >= 1).
        fn: Callable that receives a sqlite3.Connection and performs DDL/DML.
            Must NOT commit — the runner wraps everything in a transaction.
    """
    if version < 1:
        raise ValueError(f"Migration version must be >= 1, got {version}")
    _registry.setdefault(db_name, []).append((version, fn))
    # Keep sorted by version
    _registry[db_name].sort(key=lambda x: x[0])


def get_schema_version(conn: sqlite3.Connection) -> int:
    """Get current schema version for a connection.

    Returns 0 if the version table doesn't exist yet.
    """
    try:
        row = conn.execute("SELECT MAX(version) FROM _schema_version").fetchone()
        return row[0] if row and row[0] is not None else 0
    except sqlite3.OperationalError:
        return 0


def run_migrations(conn: sqlite3.Connection, db_name: str) -> int:
    """Run all pending migrations for a database.

    Args:
        conn: Open SQLite connection.
        db_name: Logical database name matching registered migrations.

    Returns:
        Number of migrations applied.
    """
    migrations = _registry.get(db_name, [])
    if not migrations:
        return 0

    # Ensure version tracking table exists
    conn.execute("""
        CREATE TABLE IF NOT EXISTS _schema_version (
            version INTEGER PRIMARY KEY,
            applied_at TEXT NOT NULL DEFAULT (datetime('now'))
        )
    """)

    current = get_schema_version(conn)
    applied = 0

    for version, fn in migrations:
        if version <= current:
            continue

        # Validate no gaps
        if version != current + 1:
            logger.error(
                "Migration gap: %s has version %d but current is %d",
                db_name,
                version,
                current,
            )
            raise RuntimeError(
                f"Migration gap for '{db_name}': expected version {current + 1}, "
                f"got {version}. Missing migration?"
            )

        logger.info("Applying migration %s v%d", db_name, version)
        try:
            fn(conn)
            conn.execute(
                "INSERT INTO _schema_version (version) VALUES (?)",
                (version,),
            )
            conn.commit()
            current = version
            applied += 1
        except Exception:
            conn.rollback()
            logger.error(
                "Migration %s v%d failed — rolled back",
                db_name,
                version,
                exc_info=True,
            )
            raise

    if applied:
        logger.info(
            "Applied %d migration(s) for %s (now at v%d)",
            applied,
            db_name,
            current,
        )

    return applied
