PicoBot/src/storage/migrations.rs

//! Schema migration helpers.
//!
//! Each `ensure_*_schema` function brings a table up to the current shape,
//! either by adding missing columns or by rebuilding the table. They run on
//! every [`super::SessionStore`] construction and are idempotent.

use rusqlite::Connection;

use super::StorageError;

pub(super) fn ensure_sessions_schema(conn: &Connection) -> Result<(), StorageError> {
    if !has_column(conn, "sessions", "user_turn_count")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE sessions ADD COLUMN user_turn_count INTEGER NOT NULL DEFAULT 0",
        )?;
    }

    if !has_column(conn, "sessions", "agent_prompt_reinjection_count")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE sessions ADD COLUMN agent_prompt_reinjection_count INTEGER NOT NULL DEFAULT 0",
        )?;
    }

    Ok(())
}

pub(super) fn ensure_messages_schema(conn: &Connection) -> Result<(), StorageError> {
    if !has_column(conn, "messages", "system_context")? {
        add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN system_context TEXT")?;
    }

    if !has_column(conn, "messages", "reasoning_content")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN reasoning_content TEXT",
        )?;
    }

    if !has_column(conn, "messages", "topic_id")? {
        add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN topic_id TEXT")?;
        // 添加外键约束（SQLite 不支持 ALTER TABLE ADD FOREIGN KEY，需要重建表）
        // 这里只添加列，外键约束由应用层保证
    }

    if !has_column(conn, "messages", "tool_duration_ms")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN tool_duration_ms INTEGER",
        )?;
    }

    // Token usage 字段（仅 assistant 消息有值，来自 LLM 响应）
    if !has_column(conn, "messages", "prompt_tokens")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN prompt_tokens INTEGER",
        )?;
    }
    if !has_column(conn, "messages", "completion_tokens")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN completion_tokens INTEGER",
        )?;
    }
    if !has_column(conn, "messages", "total_tokens")? {
        add_column_if_missing(conn, "ALTER TABLE messages ADD COLUMN total_tokens INTEGER")?;
    }
    if !has_column(conn, "messages", "context_window_tokens")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN context_window_tokens INTEGER",
        )?;
    }

    // is_compacted: 1 表示该消息是被压缩消费掉的原始消息（前端可见、LLM 不可见）。
    // 压缩摘要消息 is_compacted=0（LLM 可见），通过 system_context='history_compaction*'
    // 在前端查询中被排除。保留的原消息（system_guards / 最新 user）is_compacted=0，不重复。
    if !has_column(conn, "messages", "is_compacted")? {
        add_column_if_missing(
            conn,
            "ALTER TABLE messages ADD COLUMN is_compacted INTEGER NOT NULL DEFAULT 0",
        )?;
    }

    // 创建 topic_id 索引（如果不存在）
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_messages_topic_seq ON messages(topic_id, seq) WHERE topic_id IS NOT NULL",
        [],
    )?;

    Ok(())
}

/// topics 表：话题级模型选择列（用户在特定话题内显式选择的 provider/model）。
/// NULL 表示该话题无显式选择（运行时按 session 级 → expert → config 链解析）。
pub(super) fn ensure_topics_schema(conn: &Connection) -> Result<(), StorageError> {
    if !has_column(conn, "topics", "provider")? {
        add_column_if_missing(conn, "ALTER TABLE topics ADD COLUMN provider TEXT")?;
    }
    if !has_column(conn, "topics", "model")? {
        add_column_if_missing(conn, "ALTER TABLE topics ADD COLUMN model TEXT")?;
    }
    Ok(())
}

pub(super) fn ensure_scheduler_schema(conn: &Connection) -> Result<(), StorageError> {
    if !has_column(conn, "scheduler_jobs", "schedule_json")? {
        conn.execute(
            "ALTER TABLE scheduler_jobs ADD COLUMN schedule_json TEXT NOT NULL DEFAULT '{}'",
            [],
        )?;
    }

    if !has_column(conn, "scheduler_jobs", "state")? {
        conn.execute(
            "ALTER TABLE scheduler_jobs ADD COLUMN state TEXT NOT NULL DEFAULT 'scheduled'",
            [],
        )?;
    }

    if !has_column(conn, "scheduler_jobs", "last_status")? {
        conn.execute("ALTER TABLE scheduler_jobs ADD COLUMN last_status TEXT", [])?;
    }

    if !has_column(conn, "scheduler_jobs", "last_error")? {
        conn.execute("ALTER TABLE scheduler_jobs ADD COLUMN last_error TEXT", [])?;
    }

    if !has_column(conn, "scheduler_jobs", "run_count")? {
        conn.execute(
            "ALTER TABLE scheduler_jobs ADD COLUMN run_count INTEGER NOT NULL DEFAULT 0",
            [],
        )?;
    }

    if !has_column(conn, "scheduler_jobs", "max_runs")? {
        conn.execute("ALTER TABLE scheduler_jobs ADD COLUMN max_runs INTEGER", [])?;
    }

    if !has_column(conn, "scheduler_jobs", "paused_at")? {
        conn.execute(
            "ALTER TABLE scheduler_jobs ADD COLUMN paused_at INTEGER",
            [],
        )?;
    }

    if !has_column(conn, "scheduler_jobs", "completed_at")? {
        conn.execute(
            "ALTER TABLE scheduler_jobs ADD COLUMN completed_at INTEGER",
            [],
        )?;
    }

    Ok(())
}

pub(super) fn ensure_memory_scope_key_migration(conn: &Connection) -> Result<(), StorageError> {
    // 用 PRAGMA user_version 追踪迁移是否已完成，避免每次启动都执行全表 DELETE + UPDATE。
    // user_version 是 SQLite 内置的 32 位整数，持久化在数据库文件头中。
    // 版本 0：未迁移；版本 1：memory_scope_key 迁移已完成。
    const MEMORY_SCOPE_KEY_MIGRATION_VERSION: i64 = 1;

    let current_version: i64 = conn.query_row("PRAGMA user_version", [], |row| row.get(0))?;

    if current_version >= MEMORY_SCOPE_KEY_MIGRATION_VERSION {
        // 已迁移过，跳过
        return Ok(());
    }

    // 步骤1：去重。多条记录 scope_key 不同，改为 "default" 后会违反唯一约束。
    // 对每个 (scope_kind, namespace, memory_key) 组合保留 updated_at 最新的一条。
    conn.execute(
        "
        DELETE FROM memories
        WHERE rowid NOT IN (
            SELECT rowid FROM (
                SELECT rowid, ROW_NUMBER() OVER (
                    PARTITION BY scope_kind, namespace, memory_key
                    ORDER BY updated_at DESC
                ) AS rn
                FROM memories
            )
            WHERE rn = 1
        )
        ",
        [],
    )?;

    // 步骤2：统一 scope_key
    conn.execute(
        "UPDATE memories SET scope_key = 'default' WHERE scope_key != 'default'",
        [],
    )?;

    // 步骤3：记录迁移版本，后续启动直接跳过
    conn.execute(
        &format!(
            "PRAGMA user_version = {}",
            MEMORY_SCOPE_KEY_MIGRATION_VERSION
        ),
        [],
    )?;

    Ok(())
}

pub(super) fn ensure_todos_schema(conn: &Connection) -> Result<(), StorageError> {
    let table_exists: bool = conn
        .query_row(
            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='todos'",
            [],
            |row| row.get::<_, i64>(0),
        )
        .map(|count| count > 0)?;

    if !table_exists {
        conn.execute_batch(
            "
            CREATE TABLE IF NOT EXISTS todos (
                id TEXT NOT NULL,
                scope_key TEXT NOT NULL,
                session_id TEXT NOT NULL,
                topic_id TEXT,
                content TEXT NOT NULL,
                status TEXT NOT NULL DEFAULT 'pending',
                priority TEXT NOT NULL DEFAULT 'medium',
                created_at INTEGER NOT NULL,
                updated_at INTEGER NOT NULL,
                created_by_message_id TEXT,
                PRIMARY KEY (id, scope_key)
            );

            CREATE INDEX IF NOT EXISTS idx_todos_scope
                ON todos(scope_key, created_at ASC);

            CREATE INDEX IF NOT EXISTS idx_todos_session
                ON todos(session_id);
            ",
        )?;
        return Ok(());
    }

    // Migration: check if old schema has single-column PRIMARY KEY on `id`
    // If so, migrate to composite PRIMARY KEY (id, scope_key)
    let sql: String = conn
        .query_row(
            "SELECT sql FROM sqlite_master WHERE type='table' AND name='todos'",
            [],
            |row| row.get::<_, String>(0),
        )
        .unwrap_or_default();

    let needs_migration = sql.contains("id TEXT PRIMARY KEY")
        || (sql.contains("PRIMARY KEY") && !sql.contains("PRIMARY KEY (id, scope_key)"));

    if needs_migration {
        tracing::info!("Migrating todos table to composite PRIMARY KEY (id, scope_key)");
        conn.execute_batch(
            "
            CREATE TABLE todos_new (
                id TEXT NOT NULL,
                scope_key TEXT NOT NULL,
                session_id TEXT NOT NULL,
                topic_id TEXT,
                content TEXT NOT NULL,
                status TEXT NOT NULL DEFAULT 'pending',
                priority TEXT NOT NULL DEFAULT 'medium',
                created_at INTEGER NOT NULL,
                updated_at INTEGER NOT NULL,
                created_by_message_id TEXT,
                PRIMARY KEY (id, scope_key)
            );

            INSERT OR IGNORE INTO todos_new
                SELECT id, scope_key, session_id, topic_id, content, status, priority, created_at, updated_at
                FROM todos;

            DROP TABLE todos;

            ALTER TABLE todos_new RENAME TO todos;

            CREATE INDEX IF NOT EXISTS idx_todos_scope
                ON todos(scope_key, created_at ASC);

            CREATE INDEX IF NOT EXISTS idx_todos_session
                ON todos(session_id);
            ",
        )?;
        tracing::info!("Todos table migration complete");
    }

    // Column migration: add created_by_message_id if it doesn't exist
    let has_column = has_column(&conn, "todos", "created_by_message_id")?;
    if !has_column {
        tracing::info!("Adding created_by_message_id column to todos table");
        conn.execute(
            "ALTER TABLE todos ADD COLUMN created_by_message_id TEXT",
            [],
        )?;
        tracing::info!("Todos table column migration complete");
    }

    Ok(())
}

pub(super) fn has_column(
    conn: &Connection,
    table_name: &str,
    column_name: &str,
) -> Result<bool, StorageError> {
    let pragma = format!("PRAGMA table_info({})", table_name);
    let mut stmt = conn.prepare(&pragma)?;
    let mut rows = stmt.query([])?;

    while let Some(row) = rows.next()? {
        let existing_name: String = row.get(1)?;
        if existing_name == column_name {
            return Ok(true);
        }
    }

    Ok(false)
}

pub(super) fn add_column_if_missing(conn: &Connection, sql: &str) -> Result<(), StorageError> {
    match conn.execute(sql, []) {
        Ok(_) => Ok(()),
        Err(rusqlite::Error::SqliteFailure(_, Some(message)))
            if message.contains("duplicate column name") =>
        {
            Ok(())
        }
        Err(error) => Err(StorageError::Database(error)),
    }
}

/// pending_subagents 表：跟踪异步子代理执行状态，用于崩溃恢复。
pub(super) fn ensure_pending_subagents_schema(conn: &Connection) -> Result<(), StorageError> {
    conn.execute_batch(
        "
        CREATE TABLE IF NOT EXISTS pending_subagents (
            task_id TEXT PRIMARY KEY,
            parent_session_id TEXT NOT NULL,
            parent_topic_id TEXT NOT NULL,
            parent_chat_id TEXT NOT NULL,
            parent_channel TEXT NOT NULL,
            def_name TEXT,
            spawned_at INTEGER NOT NULL,
            status TEXT NOT NULL DEFAULT 'running'
        );

        CREATE INDEX IF NOT EXISTS idx_pending_subagents_topic
            ON pending_subagents(parent_topic_id, status);

        CREATE INDEX IF NOT EXISTS idx_pending_subagents_session
            ON pending_subagents(parent_session_id);
        ",
    )?;
    Ok(())
}