loopflow 0.10.1

use crate::lfdb::StoreResult;
use std::collections::HashSet;

#[derive(Debug, Clone)]
pub struct Migration {
    pub version: &'static str,
    pub sql: &'static str,
}

/// All migrations in order. Add new migrations here.
///
/// Each migration should work for sqlite.
const ALL_MIGRATIONS: &[Migration] = &[
    Migration {
        version: "001_initial",
        sql: include_str!("migrations/001_initial.sql"),
    },
    Migration {
        version: "002_stimulus_enabled",
        sql: include_str!("migrations/002_stimulus_enabled.sql"),
    },
    Migration {
        version: "003_agent_container_id",
        sql: include_str!("migrations/003_agent_container_id.sql"),
    },
    Migration {
        version: "004_wave_run_kind",
        sql: include_str!("migrations/004_wave_run_kind.sql"),
    },
    Migration {
        version: "005_wave_run_lineage_live_pr_state",
        sql: include_str!("migrations/005_wave_run_lineage_live_pr_state.sql"),
    },
    Migration {
        version: "006_wave_schema_provenance",
        sql: include_str!("migrations/006_wave_schema_provenance.sql"),
    },
    Migration {
        version: "007_chat_memory_blocks",
        sql: include_str!("migrations/007_chat_memory_blocks.sql"),
    },
    Migration {
        version: "008_chat_messages",
        sql: include_str!("migrations/008_chat_messages.sql"),
    },
    Migration {
        version: "009_wave_queue_state",
        sql: include_str!("migrations/009_wave_queue_state.sql"),
    },
    Migration {
        version: "011_chords_data_model",
        sql: include_str!("migrations/011_chords_data_model.sql"),
    },
    Migration {
        version: "012_drop_wave_schema_columns",
        sql: include_str!("migrations/012_drop_wave_schema_columns.sql"),
    },
    Migration {
        version: "013_remove_chord_tree",
        sql: include_str!("migrations/013_remove_chord_tree.sql"),
    },
    Migration {
        version: "015_activation_orchestration",
        sql: include_str!("migrations/015_activation_orchestration.sql"),
    },
    Migration {
        version: "016_provider_tokens",
        sql: include_str!("migrations/016_provider_tokens.sql"),
    },
    // Two migrations share the 016_ numeric prefix. The version strings stored
    // in schema_migrations are the full names, so they're distinct rows. No
    // renumbering — that would require a corrective migration.
    Migration {
        version: "016_rename_sidecar_kind_to_ci_fix_kind",
        sql: include_str!("migrations/016_rename_sidecar_kind_to_ci_fix_kind.sql"),
    },
    Migration {
        version: "017_signal_simplification",
        sql: include_str!("migrations/017_signal_simplification.sql"),
    },
    Migration {
        version: "018_wave_serialized",
        sql: include_str!("migrations/018_wave_serialized.sql"),
    },
    Migration {
        version: "019_activation_target_branch",
        sql: include_str!("migrations/019_activation_target_branch.sql"),
    },
    Migration {
        version: "020_repos",
        sql: include_str!("migrations/020_repos.sql"),
    },
    Migration {
        version: "021_repo_edges",
        sql: include_str!("migrations/021_repo_edges.sql"),
    },
    Migration {
        version: "022_stimulus_max_iterations",
        sql: include_str!("migrations/022_stimulus_max_iterations.sql"),
    },
    Migration {
        version: "023_wave_cycle_start_iteration",
        sql: include_str!("migrations/023_wave_cycle_start_iteration.sql"),
    },
    Migration {
        version: "024_signal_cleanup",
        sql: include_str!("migrations/024_signal_cleanup.sql"),
    },
    Migration {
        version: "025_credential_type",
        sql: include_str!("migrations/025_credential_type.sql"),
    },
    Migration {
        version: "026_rename_stimuli_to_triggers",
        sql: include_str!("migrations/026_rename_stimuli_to_triggers.sql"),
    },
    Migration {
        version: "027_provider_tokens_encrypted",
        sql: include_str!("migrations/027_provider_tokens_encrypted.sql"),
    },
    Migration {
        version: "028_drop_chords_tables",
        sql: include_str!("migrations/028_drop_chords_tables.sql"),
    },
    Migration {
        version: "029_attention_items",
        sql: include_str!("migrations/029_attention_items.sql"),
    },
    Migration {
        version: "030_wave_run_repair_of",
        sql: include_str!("migrations/030_wave_run_repair_of.sql"),
    },
    Migration {
        version: "031_secrets_provider",
        sql: include_str!("migrations/031_secrets_provider.sql"),
    },
    Migration {
        version: "032_sessions",
        sql: include_str!("migrations/032_terminal_sessions.sql"),
    },
    Migration {
        version: "033_wave_workers",
        sql: include_str!("migrations/033_wave_workers.sql"),
    },
    Migration {
        version: "034_wave_run_execution_cursor",
        sql: include_str!("migrations/034_wave_run_execution_cursor.sql"),
    },
    Migration {
        version: "035_session_tmux_name",
        sql: include_str!("migrations/035_terminal_session_tmux_name.sql"),
    },
    Migration {
        version: "036_wave_crons",
        sql: include_str!("migrations/036_wave_crons.sql"),
    },
    Migration {
        version: "037_wave_goal",
        sql: include_str!("migrations/037_wave_goal.sql"),
    },
    Migration {
        version: "038_wave_metrics",
        sql: include_str!("migrations/038_wave_metrics.sql"),
    },
    Migration {
        version: "039_wave_run_snapshot_task",
        sql: include_str!("migrations/039_wave_run_snapshot_task.sql"),
    },
    Migration {
        version: "040_session_use",
        sql: include_str!("migrations/040_terminal_session_use.sql"),
    },
    Migration {
        version: "041_session_parent",
        sql: include_str!("migrations/041_terminal_session_parent.sql"),
    },
    Migration {
        version: "042_wave_repos",
        sql: include_str!("migrations/042_wave_repos.sql"),
    },
    Migration {
        version: "043_drop_legacy_wave_columns",
        sql: include_str!("migrations/043_drop_legacy_wave_columns.sql"),
    },
    Migration {
        version: "044_wave_parent",
        sql: include_str!("migrations/044_wave_parent.sql"),
    },
    Migration {
        version: "045_run_token_usage",
        sql: include_str!("migrations/045_run_token_usage.sql"),
    },
    Migration {
        version: "046_run_token_usage_repo",
        sql: include_str!("migrations/046_run_token_usage_repo.sql"),
    },
    Migration {
        version: "047_run_events",
        sql: include_str!("migrations/047_run_events.sql"),
    },
    Migration {
        version: "048_terminal_sessions_run_id",
        sql: include_str!("migrations/048_terminal_sessions_run_id.sql"),
    },
    Migration {
        version: "049_runs_rename",
        sql: include_str!("migrations/049_runs_rename.sql"),
    },
    Migration {
        version: "050_drop_trigger_organs",
        sql: include_str!("migrations/050_drop_trigger_organs.sql"),
    },
    Migration {
        version: "051_drop_dead_tables",
        sql: include_str!("migrations/051_drop_dead_tables.sql"),
    },
];

/// Migrations that rename or drop schema objects some dbs never had (the
/// collapse edited historical CREATEs in place; drops may re-run against an
/// already-converged db): their "no such table/column" failure means
/// "already in the target state".
const RENAME_CONVERGENCE_MIGRATIONS: &[&str] = &[
    "048_terminal_sessions_run_id",
    "049_runs_rename",
    "050_drop_trigger_organs",
];

/// Per-migration failures that mean "the db is already in the target state":
/// record the migration as applied and converge instead of crashing.
fn is_tolerated_migration_error(version: &str, message: &str) -> bool {
    // Additive ADD COLUMN re-runs after a version-id rename.
    if message.contains("duplicate column name") || message.contains("already exists") {
        return true;
    }
    RENAME_CONVERGENCE_MIGRATIONS.contains(&version)
        && (message.contains("no such column")
            || message.contains("no such table")
            || message.contains("does not exist"))
}

pub fn migrations() -> &'static [Migration] {
    ALL_MIGRATIONS
}

// -- SQLite ------------------------------------------------------------------

pub fn apply_sqlite(conn: &rusqlite::Connection) -> StoreResult<()> {
    conn.execute_batch(
        "CREATE TABLE IF NOT EXISTS schema_migrations (
            version TEXT PRIMARY KEY,
            applied_at INTEGER NOT NULL
        )",
    )?;

    let applied = applied_versions_sqlite(conn)?;

    for migration in migrations() {
        if applied.contains(migration.version) {
            continue;
        }
        conn.execute_batch("BEGIN EXCLUSIVE")?;
        let result = (|| -> StoreResult<()> {
            // Another connection may have applied this migration between the
            // `applied` read above and taking the exclusive lock (lf and lfd
            // can open a fresh store concurrently) — re-check inside it.
            let already: bool = conn.query_row(
                "SELECT EXISTS(SELECT 1 FROM schema_migrations WHERE version = ?1)",
                rusqlite::params![migration.version],
                |row| row.get(0),
            )?;
            if already {
                conn.execute_batch("COMMIT")?;
                return Ok(());
            }
            // A migration that fails only because the db is already in its
            // target state (column already added, column already renamed) is
            // effectively applied — record it and converge rather than
            // crashing every store whose history diverged.
            match conn.execute_batch(migration.sql) {
                Ok(()) => {}
                Err(e) if is_tolerated_migration_error(migration.version, &e.to_string()) => {}
                Err(e) => return Err(e.into()),
            }
            conn.execute(
                "INSERT INTO schema_migrations (version, applied_at) VALUES (?1, ?2)",
                rusqlite::params![migration.version, now_unix()],
            )?;
            conn.execute_batch("COMMIT")?;
            Ok(())
        })();
        if let Err(e) = result {
            let _ = conn.execute_batch("ROLLBACK");
            return Err(e);
        }
    }

    Ok(())
}

fn applied_versions_sqlite(conn: &rusqlite::Connection) -> StoreResult<HashSet<String>> {
    // Table might not exist yet on first run
    let exists: bool = conn.query_row(
        "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='schema_migrations')",
        [],
        |row| row.get(0),
    )?;

    if !exists {
        return Ok(HashSet::new());
    }

    let mut stmt = conn.prepare("SELECT version FROM schema_migrations")?;
    let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
    Ok(rows.collect::<Result<HashSet<_>, _>>()?)
}

/// Latest applied migration version, or empty string if none applied.
pub fn latest_version_sqlite(conn: &rusqlite::Connection) -> StoreResult<String> {
    let applied = applied_versions_sqlite(conn)?;
    Ok(latest_applied_version(&applied))
}

fn now_unix() -> i64 {
    time::OffsetDateTime::now_utc().unix_timestamp()
}

fn latest_applied_version(applied: &HashSet<String>) -> String {
    migrations()
        .iter()
        .rev()
        .find(|m| applied.contains(m.version))
        .map(|m| m.version.to_string())
        .unwrap_or_default()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_migrations_apply_to_fresh_sqlite() {
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        apply_sqlite(&conn).unwrap();

        // All migration versions recorded
        let applied = applied_versions_sqlite(&conn).unwrap();
        let expected_count = migrations().len();
        assert_eq!(
            applied.len(),
            expected_count,
            "expected {expected_count} migrations, got {}",
            applied.len()
        );
        for migration in migrations() {
            assert!(
                applied.contains(migration.version),
                "migration {} not found in schema_migrations",
                migration.version
            );
        }

        // Key tables exist
        let tables: Vec<String> = conn
            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
            .unwrap()
            .query_map([], |row| row.get(0))
            .unwrap()
            .collect::<Result<_, _>>()
            .unwrap();
        for expected in ["waves", "repos", "runs", "terminal_sessions"] {
            assert!(
                tables.iter().any(|t| t == expected),
                "expected table {expected} not found; tables: {tables:?}"
            );
        }

        // Chords died long ago; the trigger organs died in migration 050.
        for unexpected in [
            "chords",
            "chord_members",
            "triggers",
            "stimuli",
            "pending_activations",
            "activation_log",
            "agents",
            "wave_crons",
        ] {
            assert!(
                tables.iter().all(|t| t != unexpected),
                "unexpected table {unexpected} found; tables: {tables:?}"
            );
        }

        // Re-application is idempotent
        apply_sqlite(&conn).unwrap();
        let applied_again = applied_versions_sqlite(&conn).unwrap();
        assert_eq!(applied, applied_again);
    }

    #[test]
    fn renamed_migration_id_tolerates_existing_column() {
        // Reproduces the `035` rename bug: a db that already has an additive
        // column but recorded the migration under a since-renamed id would
        // re-run its `ADD COLUMN` and crash with "duplicate column name". The
        // runner must treat that as already-applied and converge.
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        apply_sqlite(&conn).unwrap();

        // Simulate the pre-rename state: the column exists (from the first
        // apply) but the current version id is no longer recorded.
        conn.execute(
            "DELETE FROM schema_migrations WHERE version = ?1",
            rusqlite::params!["035_session_tmux_name"],
        )
        .unwrap();

        // Previously this re-ran `ALTER TABLE terminal_sessions ADD COLUMN
        // tmux_name` against a column that already exists and errored.
        apply_sqlite(&conn).expect("re-apply must tolerate the existing column");

        let applied = applied_versions_sqlite(&conn).unwrap();
        assert!(
            applied.contains("035_session_tmux_name"),
            "migration should be recorded again after convergence"
        );
    }
}

#[cfg(test)]
mod drift_tests {
    use super::*;

    #[test]
    fn rename_migration_converges_an_old_schema_db() {
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        // Old-world db: pre-collapse column name, all prior migrations recorded.
        conn.execute_batch(
            "CREATE TABLE terminal_sessions (id TEXT PRIMARY KEY, wave_id TEXT NOT NULL, wave_run_id TEXT);
             CREATE TABLE schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL);",
        )
        .unwrap();
        for m in migrations() {
            if m.version != "048_terminal_sessions_run_id" {
                conn.execute(
                    "INSERT INTO schema_migrations (version, applied_at) VALUES (?1, 0)",
                    rusqlite::params![m.version],
                )
                .unwrap();
            }
        }
        apply_sqlite(&conn).unwrap();
        let has_run_id: bool = conn
            .query_row(
                "SELECT COUNT(*) > 0 FROM pragma_table_info('terminal_sessions') WHERE name='run_id'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert!(has_run_id, "wave_run_id should be renamed to run_id");
    }

    #[test]
    fn runs_rename_converges_an_old_schema_db() {
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        // Old-world shape: pre-rename table names, and runs still carrying
        // the activation_log_id column migration 050 removes.
        conn.execute_batch(
            "CREATE TABLE wave_runs (id TEXT PRIMARY KEY, wave_id TEXT NOT NULL, activation_log_id TEXT);
             CREATE TABLE agents (id TEXT PRIMARY KEY, wave_run_id TEXT);
             CREATE TABLE fork_runs (id TEXT PRIMARY KEY, wave_run_id TEXT);
             CREATE TABLE schema_migrations (version TEXT PRIMARY KEY, applied_at INTEGER NOT NULL);",
        )
        .unwrap();
        for m in migrations() {
            if m.version != "049_runs_rename" && m.version != "050_drop_trigger_organs" {
                conn.execute(
                    "INSERT INTO schema_migrations (version, applied_at) VALUES (?1, 0)",
                    rusqlite::params![m.version],
                )
                .unwrap();
            }
        }
        apply_sqlite(&conn).unwrap();
        let runs_exists: bool = conn
            .query_row(
                "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name='runs'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert!(runs_exists, "wave_runs should be renamed to runs");
        let has_run_id: bool = conn
            .query_row(
                "SELECT COUNT(*) > 0 FROM pragma_table_info('fork_runs') WHERE name='run_id'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert!(
            has_run_id,
            "fork_runs.wave_run_id should be renamed to run_id"
        );
        // Migration 050 drops the agents table after 049 renamed its column.
        let agents_exists: bool = conn
            .query_row(
                "SELECT COUNT(*) > 0 FROM sqlite_master WHERE type='table' AND name='agents'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert!(!agents_exists, "agents should be dropped by migration 050");
    }

    #[test]
    fn rename_migration_is_tolerated_on_a_fresh_db() {
        // Fresh db: full migration chain creates run_id directly; 048's rename
        // fails benignly and must still be recorded as applied.
        let conn = rusqlite::Connection::open_in_memory().unwrap();
        apply_sqlite(&conn).unwrap();
        let recorded: bool = conn
            .query_row(
                "SELECT COUNT(*) > 0 FROM schema_migrations WHERE version='048_terminal_sessions_run_id'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        assert!(recorded);
        // Idempotent on re-run.
        apply_sqlite(&conn).unwrap();
    }
}