Lumotia/crates/storage/src/migrations.rs

use kon_core::error::{KonError, Result};
use sqlx::SqlitePool;

/// Each migration is a (version, description, sql) tuple.
/// Migrations MUST be append-only — never modify an existing migration.
/// Column defaults and NOT NULL constraints must exactly match the existing schema.
const MIGRATIONS: &[(i64, &str, &str)] = &[
    (
        1,
        "initial schema",
        r#"
        CREATE TABLE IF NOT EXISTS transcripts (
            id TEXT PRIMARY KEY,
            text TEXT NOT NULL DEFAULT '',
            source TEXT NOT NULL DEFAULT 'microphone',
            title TEXT,
            audio_path TEXT,
            duration REAL NOT NULL DEFAULT 0.0,
            engine TEXT,
            model_id TEXT,
            inference_ms INTEGER,
            sample_rate INTEGER,
            audio_channels INTEGER,
            format_mode TEXT,
            remove_fillers INTEGER NOT NULL DEFAULT 0,
            british_english INTEGER NOT NULL DEFAULT 0,
            anti_hallucination INTEGER NOT NULL DEFAULT 0,
            created_at TEXT NOT NULL DEFAULT (datetime('now'))
        );

        CREATE TABLE IF NOT EXISTS segments (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            transcript_id TEXT NOT NULL REFERENCES transcripts(id) ON DELETE CASCADE,
            start_time REAL NOT NULL,
            end_time REAL NOT NULL,
            text TEXT NOT NULL DEFAULT ''
        );

        CREATE TABLE IF NOT EXISTS tasks (
            id TEXT PRIMARY KEY,
            text TEXT NOT NULL,
            bucket TEXT NOT NULL DEFAULT 'inbox',
            list_id TEXT,
            effort TEXT,
            done INTEGER NOT NULL DEFAULT 0,
            done_at TEXT,
            created_at TEXT NOT NULL DEFAULT (datetime('now')),
            source_transcript_id TEXT
        );

        CREATE TABLE IF NOT EXISTS task_lists (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL,
            built_in INTEGER NOT NULL DEFAULT 0,
            profile_id TEXT,
            created_at TEXT NOT NULL DEFAULT (datetime('now'))
        );

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

        CREATE TABLE IF NOT EXISTS error_log (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT NOT NULL DEFAULT (datetime('now')),
            context TEXT NOT NULL,
            error_code TEXT,
            message TEXT NOT NULL,
            metadata TEXT
        );

        CREATE INDEX IF NOT EXISTS idx_transcripts_created ON transcripts(created_at);
        CREATE INDEX IF NOT EXISTS idx_segments_transcript ON segments(transcript_id);
        CREATE INDEX IF NOT EXISTS idx_tasks_bucket ON tasks(bucket);
        CREATE INDEX IF NOT EXISTS idx_tasks_transcript ON tasks(source_transcript_id);
        CREATE INDEX IF NOT EXISTS idx_error_log_context ON error_log(context)
    "#,
    ),
    (
        2,
        "transcripts FTS5 + dictionary table",
        r#"
        CREATE VIRTUAL TABLE IF NOT EXISTS transcripts_fts USING fts5(
            text,
            title,
            content='transcripts',
            content_rowid='rowid',
            tokenize='porter unicode61 remove_diacritics 2'
        );

        CREATE TRIGGER IF NOT EXISTS transcripts_ai AFTER INSERT ON transcripts BEGIN
            INSERT INTO transcripts_fts(rowid, text, title)
            VALUES (new.rowid, new.text, COALESCE(new.title, ''));
        END;

        CREATE TRIGGER IF NOT EXISTS transcripts_ad AFTER DELETE ON transcripts BEGIN
            INSERT INTO transcripts_fts(transcripts_fts, rowid, text, title)
            VALUES ('delete', old.rowid, old.text, COALESCE(old.title, ''));
        END;

        CREATE TRIGGER IF NOT EXISTS transcripts_au AFTER UPDATE ON transcripts BEGIN
            INSERT INTO transcripts_fts(transcripts_fts, rowid, text, title)
            VALUES ('delete', old.rowid, old.text, COALESCE(old.title, ''));
            INSERT INTO transcripts_fts(rowid, text, title)
            VALUES (new.rowid, new.text, COALESCE(new.title, ''));
        END;

        CREATE TABLE IF NOT EXISTS dictionary (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            term TEXT NOT NULL UNIQUE,
            note TEXT,
            created_at TEXT NOT NULL DEFAULT (datetime('now'))
        );

        CREATE INDEX IF NOT EXISTS idx_dictionary_term ON dictionary(term)
    "#,
    ),
    (
        3,
        "micro-stepping: parent_task_id on tasks",
        r#"
        ALTER TABLE tasks ADD COLUMN parent_task_id TEXT REFERENCES tasks(id) ON DELETE CASCADE;
        CREATE INDEX IF NOT EXISTS idx_tasks_parent ON tasks(parent_task_id)
    "#,
    ),
    (
        4,
        "tasks_meta: notes column for persisted free-text",
        r#"
        ALTER TABLE tasks ADD COLUMN notes TEXT NOT NULL DEFAULT ''
    "#,
    ),
    (
        5,
        "transcripts_meta",
        r#"
        ALTER TABLE transcripts ADD COLUMN starred INTEGER NOT NULL DEFAULT 0;
        ALTER TABLE transcripts ADD COLUMN manual_tags TEXT NOT NULL DEFAULT '';
        ALTER TABLE transcripts ADD COLUMN template TEXT NOT NULL DEFAULT '';
        ALTER TABLE transcripts ADD COLUMN language TEXT NOT NULL DEFAULT '';
        ALTER TABLE transcripts ADD COLUMN segments_json TEXT NOT NULL DEFAULT '';
        "#,
    ),
    (
        6,
        "profiles",
        r#"
        CREATE TABLE IF NOT EXISTS profiles (
            id TEXT PRIMARY KEY,
            name TEXT NOT NULL UNIQUE,
            initial_prompt TEXT NOT NULL DEFAULT '',
            created_at TEXT NOT NULL
        );

        CREATE TABLE IF NOT EXISTS profile_terms (
            id TEXT PRIMARY KEY,
            profile_id TEXT NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,
            term TEXT NOT NULL,
            note TEXT NOT NULL DEFAULT '',
            created_at TEXT NOT NULL,
            UNIQUE(profile_id, term)
        );

        CREATE INDEX IF NOT EXISTS idx_profile_terms_profile_id
            ON profile_terms(profile_id);

        INSERT OR IGNORE INTO profiles (id, name, initial_prompt, created_at)
        VALUES ('00000000-0000-0000-0000-000000000001', 'Default', '', datetime('now'));

        INSERT OR IGNORE INTO profile_terms (id, profile_id, term, note, created_at)
        SELECT
            d.id,
            '00000000-0000-0000-0000-000000000001',
            d.term,
            COALESCE(d.note, ''),
            d.created_at
        FROM dictionary d;

        CREATE TRIGGER IF NOT EXISTS trg_protect_default_profile_delete
        BEFORE DELETE ON profiles
        FOR EACH ROW
        WHEN OLD.id = '00000000-0000-0000-0000-000000000001'
        BEGIN
            SELECT RAISE(ABORT, 'Default profile cannot be deleted');
        END;

        CREATE TRIGGER IF NOT EXISTS trg_protect_default_profile_rename
        BEFORE UPDATE OF name ON profiles
        FOR EACH ROW
        WHEN OLD.id = '00000000-0000-0000-0000-000000000001'
          AND NEW.name != OLD.name
        BEGIN
            SELECT RAISE(ABORT, 'Default profile cannot be renamed');
        END;
        "#,
    ),
    (
        7,
        "drop_dictionary",
        r#"
        DROP INDEX IF EXISTS idx_dictionary_term;
        DROP TABLE IF EXISTS dictionary;
        "#,
    ),
    (
        8,
        "transcript_profile_provenance",
        r#"
        ALTER TABLE transcripts
            ADD COLUMN profile_id TEXT NOT NULL
            DEFAULT '00000000-0000-0000-0000-000000000001';

        CREATE INDEX IF NOT EXISTS idx_transcripts_profile_id
            ON transcripts(profile_id);
        "#,
    ),
];

/// Split SQL into individual statements, respecting BEGIN...END trigger blocks.
fn split_statements(sql: &str) -> Vec<String> {
    let mut result = Vec::new();
    let mut current = String::new();
    let mut depth: i32 = 0;
    let upper = sql.to_ascii_uppercase();
    let ub = upper.as_bytes();
    let ob = sql.as_bytes();
    let n = ob.len();
    let mut i = 0;
    while i < n {
        let prev_alpha = i > 0 && ob[i - 1].is_ascii_alphanumeric();
        if !prev_alpha && ub[i..].starts_with(b"BEGIN") {
            let next_alpha = i + 5 < n && ob[i + 5].is_ascii_alphanumeric();
            if !next_alpha {
                depth += 1;
            }
        }
        if !prev_alpha && ub[i..].starts_with(b"END") {
            let next_alpha = i + 3 < n && ob[i + 3].is_ascii_alphanumeric();
            if !next_alpha && depth > 0 {
                depth -= 1;
            }
        }
        if ob[i] == b';' && depth == 0 {
            let stmt = current.trim().to_string();
            if !stmt.is_empty() {
                result.push(stmt);
            }
            current = String::new();
        } else {
            current.push(ob[i] as char);
        }
        i += 1;
    }
    let stmt = current.trim().to_string();
    if !stmt.is_empty() {
        result.push(stmt);
    }
    result
}

/// Ensure the schema_version table exists and run any pending migrations.
pub async fn run_migrations(pool: &SqlitePool) -> Result<()> {
    sqlx::query(
        "CREATE TABLE IF NOT EXISTS schema_version (
            version INTEGER PRIMARY KEY,
            description TEXT NOT NULL,
            applied_at TEXT NOT NULL DEFAULT (datetime('now'))
        )",
    )
    .execute(pool)
    .await
    .map_err(|e| KonError::StorageError(format!("Schema version table creation failed: {e}")))?;

    let current: i64 = sqlx::query_scalar("SELECT COALESCE(MAX(version), 0) FROM schema_version")
        .fetch_one(pool)
        .await
        .map_err(|e| KonError::StorageError(format!("Schema version query failed: {e}")))?;

    for (version, description, sql) in MIGRATIONS {
        if *version > current {
            log::info!("Running migration {}: {}", version, description);

            let statements = split_statements(sql);

            for statement in &statements {
                sqlx::query(statement).execute(pool).await.map_err(|e| {
                    KonError::StorageError(format!("Migration {} failed: {e}", version))
                })?;
            }

            sqlx::query("INSERT INTO schema_version (version, description) VALUES (?, ?)")
                .bind(version)
                .bind(description)
                .execute(pool)
                .await
                .map_err(|e| {
                    KonError::StorageError(format!("Migration version record failed: {e}"))
                })?;

            log::info!("Migration {} complete", version);
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use sqlx::sqlite::SqlitePoolOptions;
    use sqlx::Row;

    #[tokio::test]
    async fn test_migrations_run_on_empty_db() {
        let pool = SqlitePoolOptions::new()
            .connect("sqlite::memory:")
            .await
            .unwrap();

        run_migrations(&pool).await.unwrap();

        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM schema_version")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(count, 8);

        sqlx::query("INSERT INTO settings (key, value) VALUES ('test', 'value')")
            .execute(&pool)
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn test_migrations_idempotent() {
        let pool = SqlitePoolOptions::new()
            .connect("sqlite::memory:")
            .await
            .unwrap();

        run_migrations(&pool).await.unwrap();
        run_migrations(&pool).await.unwrap();

        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM schema_version")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(count, 8);
    }

    #[tokio::test]
    async fn migration_tasks_meta_adds_columns() {
        // Task 2.6 — verify list_id / effort / notes are all present on the
        // tasks table after migrations run. list_id and effort have been
        // present since v1 (nullable); notes is added by v4.
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let info = sqlx::query("PRAGMA table_info(tasks)")
            .fetch_all(&pool)
            .await
            .unwrap();
        let names: Vec<String> = info.iter().map(|r| r.get::<String, _>("name")).collect();
        for col in ["list_id", "effort", "notes"] {
            assert!(
                names.contains(&col.to_string()),
                "tasks must have {col}; got {names:?}"
            );
        }
    }

    #[tokio::test]
    async fn migration_transcripts_meta_adds_columns() {
        // Task 2.5 — verify starred / manual_tags / template / language /
        // segments_json are all present on the transcripts table after
        // migrations run. All five are added by v5.
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let info = sqlx::query("PRAGMA table_info(transcripts)")
            .fetch_all(&pool)
            .await
            .unwrap();
        let names: Vec<String> = info.iter().map(|r| r.get::<String, _>("name")).collect();
        for col in [
            "starred",
            "manual_tags",
            "template",
            "language",
            "segments_json",
        ] {
            assert!(
                names.contains(&col.to_string()),
                "transcripts must have {col}; got {names:?}"
            );
        }
    }

    #[tokio::test]
    async fn migration_transcript_profile_provenance_adds_profile_id() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let info = sqlx::query("PRAGMA table_info(transcripts)")
            .fetch_all(&pool)
            .await
            .unwrap();
        let names: Vec<String> = info.iter().map(|r| r.get::<String, _>("name")).collect();
        assert!(
            names.contains(&"profile_id".to_string()),
            "transcripts must have profile_id; got {names:?}"
        );
    }

    #[tokio::test]
    async fn test_parent_task_id_cascade_delete() {
        let pool = SqlitePoolOptions::new()
            .connect("sqlite::memory:")
            .await
            .unwrap();

        run_migrations(&pool).await.unwrap();

        // Insert parent task
        sqlx::query("INSERT INTO tasks (id, text) VALUES ('parent-1', 'Parent task')")
            .execute(&pool)
            .await
            .unwrap();

        // Insert child task with parent_task_id
        sqlx::query("INSERT INTO tasks (id, text, parent_task_id) VALUES ('child-1', 'Child task', 'parent-1')")
            .execute(&pool)
            .await
            .unwrap();

        // Delete parent — child should cascade
        sqlx::query("DELETE FROM tasks WHERE id = 'parent-1'")
            .execute(&pool)
            .await
            .unwrap();

        let remaining: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM tasks")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(
            remaining, 0,
            "cascade delete should remove child when parent is deleted"
        );
    }

    /// Test-only helper: run migrations only up to (and including) `target_version`.
    /// Mirrors `run_migrations` but stops early — used by the v6 upgrade-path test
    /// to seed a v5 schema with dictionary rows before applying v6.
    async fn run_migrations_up_to(pool: &SqlitePool, target_version: i64) -> Result<()> {
        sqlx::query(
            "CREATE TABLE IF NOT EXISTS schema_version (
                version INTEGER PRIMARY KEY,
                description TEXT NOT NULL,
                applied_at TEXT NOT NULL DEFAULT (datetime('now'))
            )",
        )
        .execute(pool)
        .await
        .map_err(|e| {
            KonError::StorageError(format!("Schema version table creation failed: {e}"))
        })?;

        let current: i64 =
            sqlx::query_scalar("SELECT COALESCE(MAX(version), 0) FROM schema_version")
                .fetch_one(pool)
                .await
                .map_err(|e| KonError::StorageError(format!("Schema version query failed: {e}")))?;

        for (version, description, sql) in MIGRATIONS {
            if *version > current && *version <= target_version {
                let statements = split_statements(sql);
                for statement in &statements {
                    sqlx::query(statement).execute(pool).await.map_err(|e| {
                        KonError::StorageError(format!("Migration {} failed: {e}", version))
                    })?;
                }
                sqlx::query("INSERT INTO schema_version (version, description) VALUES (?, ?)")
                    .bind(version)
                    .bind(description)
                    .execute(pool)
                    .await
                    .map_err(|e| {
                        KonError::StorageError(format!("Migration version record failed: {e}"))
                    })?;
            }
        }
        Ok(())
    }

    #[tokio::test]
    async fn migration_v6_seeds_default_profile_on_fresh_db() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM profiles")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(count, 1, "Default profile must be seeded on fresh install");

        let name: String = sqlx::query_scalar("SELECT name FROM profiles WHERE id = ?")
            .bind(crate::DEFAULT_PROFILE_ID)
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(name, "Default");

        let terms_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM profile_terms")
            .fetch_one(&pool)
            .await
            .unwrap();
        assert_eq!(terms_count, 0, "Fresh DB has no dictionary rows to copy");
    }

    #[tokio::test]
    async fn migration_v6_copies_dictionary_rows_to_default_profile_terms() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");

        // Bring schema up to v5 (dictionary + transcripts_meta, no profile_terms).
        run_migrations_up_to(&pool, 5).await.expect("migrate to v5");

        // dictionary.id is INTEGER PK AUTOINCREMENT (see v2); let SQLite assign rowids.
        sqlx::query(
            "INSERT INTO dictionary (term, note, created_at) VALUES \
             ('Kon', '', datetime('now')), \
             ('CORBEL', 'brand', datetime('now')), \
             ('Wren', '', datetime('now'))",
        )
        .execute(&pool)
        .await
        .expect("seed dictionary");

        run_migrations(&pool).await.expect("migrate to v6");

        let copied: i64 =
            sqlx::query_scalar("SELECT COUNT(*) FROM profile_terms WHERE profile_id = ?")
                .bind(crate::DEFAULT_PROFILE_ID)
                .fetch_one(&pool)
                .await
                .unwrap();
        assert_eq!(copied, 3);

        let corbel_note: String =
            sqlx::query_scalar("SELECT note FROM profile_terms WHERE term = 'CORBEL'")
                .fetch_one(&pool)
                .await
                .unwrap();
        assert_eq!(corbel_note, "brand");
    }

    #[tokio::test]
    async fn migration_v6_trigger_rejects_default_profile_delete() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let result = sqlx::query("DELETE FROM profiles WHERE id = ?")
            .bind(crate::DEFAULT_PROFILE_ID)
            .execute(&pool)
            .await;

        assert!(result.is_err(), "trigger must block Default deletion");
        let msg = result.unwrap_err().to_string();
        assert!(
            msg.contains("Default profile cannot be deleted"),
            "got: {msg}"
        );
    }

    #[tokio::test]
    async fn migration_v6_trigger_rejects_default_profile_rename() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        let result = sqlx::query("UPDATE profiles SET name = 'NotDefault' WHERE id = ?")
            .bind(crate::DEFAULT_PROFILE_ID)
            .execute(&pool)
            .await;

        assert!(result.is_err(), "trigger must block Default rename");
    }

    #[tokio::test]
    async fn migration_v7_drops_dictionary_table() {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .expect("pool");
        run_migrations(&pool).await.expect("migrate");

        // dictionary table should not exist after v7.
        let result: std::result::Result<i64, sqlx::Error> =
            sqlx::query_scalar("SELECT COUNT(*) FROM dictionary")
                .fetch_one(&pool)
                .await;
        assert!(result.is_err(), "dictionary table must be gone after v7");
        let err = result.unwrap_err().to_string().to_lowercase();
        assert!(err.contains("no such table"), "got: {err}");
    }
}