mixvideo-v2/apps/desktop/src-tauri/src/infrastructure/database/migrations.rs

use rusqlite::Connection;
use anyhow::{Result, anyhow};
use std::collections::HashMap;

/// 数据库迁移版本号
pub type MigrationVersion = u32;

/// 单个迁移脚本
#[derive(Debug, Clone)]
pub struct Migration {
    /// 迁移版本号
    pub version: MigrationVersion,
    /// 迁移描述
    pub description: String,
    /// 向前迁移SQL
    pub up_sql: String,
    /// 回滚迁移SQL（可选）
    pub down_sql: Option<String>,
}

/// 迁移管理器
pub struct MigrationManager {
    migrations: HashMap<MigrationVersion, Migration>,
}

impl MigrationManager {
    /// 创建新的迁移管理器
    pub fn new() -> Self {
        let mut manager = Self {
            migrations: HashMap::new(),
        };
        
        // 注册所有迁移
        manager.register_migrations();
        manager
    }

    /// 注册所有迁移脚本
    fn register_migrations(&mut self) {
        // 迁移 1: 初始化数据库表结构
        self.add_migration(Migration {
            version: 1,
            description: "初始化数据库表结构".to_string(),
            up_sql: include_str!("migrations/001_initial_schema.sql").to_string(),
            down_sql: Some(include_str!("migrations/001_initial_schema_down.sql").to_string()),
        });

        // 迁移 2: 添加模特ID字段到素材表
        self.add_migration(Migration {
            version: 2,
            description: "添加模特ID字段到素材表".to_string(),
            up_sql: include_str!("migrations/002_add_model_id_to_materials.sql").to_string(),
            down_sql: Some(include_str!("migrations/002_add_model_id_to_materials_down.sql").to_string()),
        });

        // 迁移 3: 添加文件存在状态字段到模板素材表
        self.add_migration(Migration {
            version: 3,
            description: "添加文件存在状态字段到模板素材表".to_string(),
            up_sql: include_str!("migrations/003_add_file_status_to_template_materials.sql").to_string(),
            down_sql: Some(include_str!("migrations/003_add_file_status_to_template_materials_down.sql").to_string()),
        });

        // 迁移 4: 移除模板表的project_id字段
        self.add_migration(Migration {
            version: 4,
            description: "移除模板表的project_id字段".to_string(),
            up_sql: include_str!("migrations/004_remove_project_id_from_templates.sql").to_string(),
            down_sql: None, // 这个迁移不支持回滚，因为涉及数据丢失
        });

        // 迁移 5: 修复AI分类表的is_active字段类型
        self.add_migration(Migration {
            version: 5,
            description: "修复AI分类表的is_active字段类型".to_string(),
            up_sql: include_str!("migrations/005_fix_ai_classifications_is_active.sql").to_string(),
            down_sql: None, // 类型修复不支持回滚
        });

        // 迁移 6: 添加导出状态字段到模板匹配结果表
        self.add_migration(Migration {
            version: 6,
            description: "添加导出状态字段到模板匹配结果表".to_string(),
            up_sql: include_str!("migrations/006_add_export_status_to_template_matching.sql").to_string(),
            down_sql: Some(include_str!("migrations/006_add_export_status_to_template_matching_down.sql").to_string()),
        });

        // 迁移 7: 修复项目表的is_active字段数据
        self.add_migration(Migration {
            version: 7,
            description: "修复项目表的is_active字段数据".to_string(),
            up_sql: include_str!("migrations/007_fix_projects_is_active_data.sql").to_string(),
            down_sql: None, // 数据修复不支持回滚
        });

        // 迁移 8: 添加缩略图和使用跟踪字段
        self.add_migration(Migration {
            version: 8,
            description: "添加缩略图和使用跟踪字段".to_string(),
            up_sql: include_str!("migrations/008_add_thumbnail_and_usage_tracking.sql").to_string(),
            down_sql: Some(include_str!("migrations/008_add_thumbnail_and_usage_tracking_down.sql").to_string()),
        });

        // 迁移 9: 修复template_materials表的file_size字段，允许为NULL
        self.add_migration(Migration {
            version: 9,
            description: "修复template_materials表的file_size字段，允许为NULL".to_string(),
            up_sql: include_str!("migrations/009_fix_template_materials_file_size_nullable.sql").to_string(),
            down_sql: Some(include_str!("migrations/009_fix_template_materials_file_size_nullable_down.sql").to_string()),
        });

        // 迁移 10: 修复video_classification_records表结构，添加缺失字段
        self.add_migration(Migration {
            version: 10,
            description: "修复video_classification_records表结构，添加缺失字段".to_string(),
            up_sql: include_str!("migrations/010_fix_video_classification_records_schema.sql").to_string(),
            down_sql: Some(include_str!("migrations/010_fix_video_classification_records_schema_down.sql").to_string()),
        });

        // 迁移 11: 修复template_matching_results表结构，添加缺失字段
        self.add_migration(Migration {
            version: 11,
            description: "修复template_matching_results表结构，添加缺失字段".to_string(),
            up_sql: include_str!("migrations/011_fix_template_matching_results_schema.sql").to_string(),
            down_sql: Some(include_str!("migrations/011_fix_template_matching_results_schema_down.sql").to_string()),
        });

        // 迁移 12: 创建video_classification_tasks表
        self.add_migration(Migration {
            version: 12,
            description: "创建video_classification_tasks表".to_string(),
            up_sql: include_str!("migrations/012_create_video_classification_tasks_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/012_create_video_classification_tasks_table_down.sql").to_string()),
        });

        // 迁移 13: 创建export_records表
        self.add_migration(Migration {
            version: 13,
            description: "创建export_records表".to_string(),
            up_sql: include_str!("migrations/013_create_export_records_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/013_create_export_records_table_down.sql").to_string()),
        });

        // 迁移 14: 创建project_template_bindings表
        self.add_migration(Migration {
            version: 14,
            description: "创建project_template_bindings表".to_string(),
            up_sql: include_str!("migrations/014_create_project_template_bindings_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/014_create_project_template_bindings_table_down.sql").to_string()),
        });

        // 迁移 15: 添加默认的AI分类数据
        self.add_migration(Migration {
            version: 15,
            description: "添加默认的AI分类数据".to_string(),
            up_sql: include_str!("migrations/015_add_default_ai_classifications.sql").to_string(),
            down_sql: Some(include_str!("migrations/015_add_default_ai_classifications_down.sql").to_string()),
        });

        // 迁移 16: 修复AI分类表中的日期时间格式
        self.add_migration(Migration {
            version: 16,
            description: "修复AI分类表中的日期时间格式".to_string(),
            up_sql: include_str!("migrations/016_fix_ai_classifications_datetime_format.sql").to_string(),
            down_sql: Some(include_str!("migrations/016_fix_ai_classifications_datetime_format_down.sql").to_string()),
        });

        // 迁移 17: 创建matching_segment_results和matching_failed_segment_results表
        self.add_migration(Migration {
            version: 17,
            description: "创建matching_segment_results和matching_failed_segment_results表".to_string(),
            up_sql: include_str!("migrations/017_create_matching_segment_results_tables.sql").to_string(),
            down_sql: Some(include_str!("migrations/017_create_matching_segment_results_tables_down.sql").to_string()),
        });

        // 迁移 18: 创建material_usage_records表
        self.add_migration(Migration {
            version: 18,
            description: "创建material_usage_records表".to_string(),
            up_sql: include_str!("migrations/018_create_material_usage_records_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/018_create_material_usage_records_table_down.sql").to_string()),
        });

        // 迁移 19: 为ai_classifications表添加weight字段
        self.add_migration(Migration {
            version: 19,
            description: "为ai_classifications表添加weight字段".to_string(),
            up_sql: include_str!("migrations/019_add_weight_to_ai_classifications.sql").to_string(),
            down_sql: Some(include_str!("migrations/019_add_weight_to_ai_classifications_down.sql").to_string()),
        });

        // 迁移 20: 创建模板片段权重配置表
        self.add_migration(Migration {
            version: 20,
            description: "创建模板片段权重配置表".to_string(),
            up_sql: include_str!("migrations/020_create_template_segment_weights_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/020_create_template_segment_weights_table_down.sql").to_string()),
        });

        // 迁移 21: 创建穿搭方案收藏表
        self.add_migration(Migration {
            version: 21,
            description: "创建穿搭方案收藏表".to_string(),
            up_sql: include_str!("migrations/021_outfit_favorites.sql").to_string(),
            down_sql: Some(include_str!("migrations/021_outfit_favorites_down.sql").to_string()),
        });

        // 迁移 22: 创建图片生成记录表
        self.add_migration(Migration {
            version: 22,
            description: "创建图片生成记录表".to_string(),
            up_sql: include_str!("migrations/022_create_image_generation_records_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/022_create_image_generation_records_table_down.sql").to_string()),
        });

        // 迁移 23: 创建穿搭图片相关表
        self.add_migration(Migration {
            version: 23,
            description: "创建穿搭图片相关表".to_string(),
            up_sql: include_str!("migrations/023_create_outfit_images_tables.sql").to_string(),
            down_sql: Some(include_str!("migrations/023_create_outfit_images_tables_down.sql").to_string()),
        });

        // 迁移 23: 创建声音克隆记录表
        self.add_migration(Migration {
            version: 23,
            description: "创建声音克隆记录表".to_string(),
            up_sql: include_str!("migrations/023_create_voice_clone_table.sql").to_string(),
            down_sql: Some(include_str!("migrations/023_create_voice_clone_table_down.sql").to_string()),
        });

        // 迁移 24: 创建语音生成记录表
        self.add_migration(Migration {
            version: 24,
            description: "创建语音生成记录表".to_string(),
            up_sql: include_str!("migrations/024_create_speech_generation_records_table.sql").to_string(),
            down_sql: None, // 暂时不提供回滚脚本
        });
    }

    /// 添加迁移
    fn add_migration(&mut self, migration: Migration) {
        self.migrations.insert(migration.version, migration);
    }

    /// 获取当前数据库版本
    pub fn get_current_version(&self, conn: &Connection) -> Result<MigrationVersion> {
        // 确保迁移历史表存在
        self.ensure_migration_table(conn)?;

        let version = conn.query_row(
            "SELECT MAX(version) FROM schema_migrations WHERE success = 1",
            [],
            |row| row.get::<_, Option<u32>>(0)
        ).unwrap_or(Some(0)).unwrap_or(0);

        Ok(version)
    }

    /// 获取最新版本号
    pub fn get_latest_version(&self) -> MigrationVersion {
        self.migrations.keys().max().copied().unwrap_or(0)
    }

    /// 确保迁移历史表存在
    fn ensure_migration_table(&self, conn: &Connection) -> Result<()> {
        conn.execute(
            "CREATE TABLE IF NOT EXISTS schema_migrations (
                version INTEGER PRIMARY KEY,
                description TEXT NOT NULL,
                applied_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                success INTEGER NOT NULL DEFAULT 0,
                error_message TEXT
            )",
            [],
        )?;
        Ok(())
    }

    /// 运行所有待执行的迁移
    pub fn migrate(&self, conn: &Connection) -> Result<()> {
        let current_version = self.get_current_version(conn)?;
        let latest_version = self.get_latest_version();

        if current_version >= latest_version {
            println!("数据库已是最新版本 v{}", current_version);
            return Ok(());
        }

        println!("开始数据库迁移: v{} -> v{}", current_version, latest_version);

        // 按版本号排序执行迁移
        let mut versions: Vec<_> = self.migrations.keys().collect();
        versions.sort();

        for &version in versions {
            if version > current_version {
                self.apply_migration(conn, version)?;
            }
        }

        println!("数据库迁移完成");
        Ok(())
    }

    /// 应用单个迁移
    fn apply_migration(&self, conn: &Connection, version: MigrationVersion) -> Result<()> {
        let migration = self.migrations.get(&version)
            .ok_or_else(|| anyhow!("迁移版本 {} 不存在", version))?;

        println!("应用迁移 v{}: {}", version, migration.description);

        // 开始事务
        let tx = conn.unchecked_transaction()?;

        // 记录迁移开始
        tx.execute(
            "INSERT OR REPLACE INTO schema_migrations (version, description, success) VALUES (?1, ?2, 0)",
            [&version.to_string(), &migration.description],
        )?;

        // 执行迁移SQL
        match self.execute_migration_sql(&tx, &migration.up_sql) {
            Ok(_) => {
                // 标记迁移成功
                tx.execute(
                    "UPDATE schema_migrations SET success = 1, applied_at = CURRENT_TIMESTAMP WHERE version = ?1",
                    [&version.to_string()],
                )?;
                tx.commit()?;
                println!("迁移 v{} 应用成功", version);
                Ok(())
            }
            Err(e) => {
                // 记录错误信息
                let _ = tx.execute(
                    "UPDATE schema_migrations SET error_message = ?1 WHERE version = ?2",
                    [&e.to_string(), &version.to_string()],
                );
                tx.rollback()?;
                Err(anyhow!("迁移 v{} 失败: {}", version, e))
            }
        }
    }

    /// 执行迁移SQL（支持多条语句）
    fn execute_migration_sql(&self, conn: &Connection, sql: &str) -> Result<()> {
        // 使用SQLite的execute_batch方法来执行多条SQL语句
        // 这个方法能正确处理SQL语句的分割，包括注释和字符串中的分号
        conn.execute_batch(sql)?;
        Ok(())
    }

    /// 回滚到指定版本
    pub fn rollback_to(&self, conn: &Connection, target_version: MigrationVersion) -> Result<()> {
        let current_version = self.get_current_version(conn)?;

        if target_version >= current_version {
            return Err(anyhow!("目标版本 {} 不能大于等于当前版本 {}", target_version, current_version));
        }

        println!("开始回滚数据库: v{} -> v{}", current_version, target_version);

        // 按版本号倒序回滚
        let mut versions: Vec<_> = self.migrations.keys().collect();
        versions.sort_by(|a, b| b.cmp(a)); // 倒序

        for &version in versions {
            if version > target_version && version <= current_version {
                self.rollback_migration(conn, version)?;
            }
        }

        println!("数据库回滚完成");
        Ok(())
    }

    /// 回滚单个迁移
    fn rollback_migration(&self, conn: &Connection, version: MigrationVersion) -> Result<()> {
        let migration = self.migrations.get(&version)
            .ok_or_else(|| anyhow!("迁移版本 {} 不存在", version))?;

        let down_sql = migration.down_sql.as_ref()
            .ok_or_else(|| anyhow!("迁移 v{} 不支持回滚", version))?;

        println!("回滚迁移 v{}: {}", version, migration.description);

        // 开始事务
        let tx = conn.unchecked_transaction()?;

        // 执行回滚SQL
        match self.execute_migration_sql(&tx, down_sql) {
            Ok(_) => {
                // 删除迁移记录
                tx.execute(
                    "DELETE FROM schema_migrations WHERE version = ?1",
                    [&version.to_string()],
                )?;
                tx.commit()?;
                println!("迁移 v{} 回滚成功", version);
                Ok(())
            }
            Err(e) => {
                tx.rollback()?;
                Err(anyhow!("迁移 v{} 回滚失败: {}", version, e))
            }
        }
    }

    /// 获取迁移历史
    pub fn get_migration_history(&self, conn: &Connection) -> Result<Vec<(MigrationVersion, String, String, bool)>> {
        self.ensure_migration_table(conn)?;

        let mut stmt = conn.prepare(
            "SELECT version, description, applied_at, success FROM schema_migrations ORDER BY version"
        )?;

        let rows = stmt.query_map([], |row| {
            Ok((
                row.get::<_, u32>(0)?,
                row.get::<_, String>(1)?,
                row.get::<_, String>(2)?,
                row.get::<_, i32>(3)? == 1,
            ))
        })?;

        let mut history = Vec::new();
        for row in rows {
            history.push(row?);
        }

        Ok(history)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rusqlite::Connection;
    use tempfile::NamedTempFile;

    fn create_test_connection() -> Connection {
        let temp_file = NamedTempFile::new().unwrap();
        let conn = Connection::open(temp_file.path()).unwrap();

        // 配置数据库设置
        conn.pragma_update(None, "foreign_keys", "ON").unwrap();
        conn.pragma_update(None, "journal_mode", "WAL").unwrap();

        conn
    }

    #[test]
    fn test_migration_manager_creation() {
        let manager = MigrationManager::new();
        let latest_version = manager.get_latest_version();

        // 应该有至少8个迁移版本
        assert!(latest_version >= 8, "应该有至少8个迁移版本，实际: {}", latest_version);
    }

    #[test]
    fn test_initial_version_is_zero() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        let current_version = manager.get_current_version(&conn).unwrap();
        assert_eq!(current_version, 0, "新数据库的版本应该是0");
    }

    #[test]
    fn test_migration_table_creation() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        // 确保迁移表存在
        manager.ensure_migration_table(&conn).unwrap();

        // 验证表是否存在
        let table_exists: bool = conn.query_row(
            "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='schema_migrations'",
            [],
            |row| row.get::<_, i32>(0)
        ).unwrap() > 0;

        assert!(table_exists, "schema_migrations表应该存在");
    }

    #[test]
    fn test_full_migration() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        // 执行所有迁移
        manager.migrate(&conn).unwrap();

        // 验证版本已更新
        let current_version = manager.get_current_version(&conn).unwrap();
        let latest_version = manager.get_latest_version();

        assert_eq!(current_version, latest_version, "迁移后版本应该是最新版本");

        // 验证一些关键表是否存在
        let tables = vec![
            "projects", "materials", "material_segments", "models",
            "ai_classifications", "templates", "template_materials",
            "tracks", "track_segments", "video_classification_records"
        ];

        for table in tables {
            let table_exists: bool = conn.query_row(
                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
                [table],
                |row| row.get::<_, i32>(0)
            ).unwrap() > 0;

            assert!(table_exists, "表 {} 应该存在", table);
        }
    }

    #[test]
    fn test_migration_history() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        // 执行迁移
        manager.migrate(&conn).unwrap();

        // 获取迁移历史
        let history = manager.get_migration_history(&conn).unwrap();

        // 应该有迁移记录
        assert!(!history.is_empty(), "应该有迁移历史记录");

        // 所有迁移都应该成功
        for (version, description, _applied_at, success) in history {
            assert!(success, "迁移 v{}: {} 应该成功", version, description);
        }
    }

    #[test]
    fn test_idempotent_migration() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        // 第一次迁移
        manager.migrate(&conn).unwrap();
        let version_after_first = manager.get_current_version(&conn).unwrap();

        // 第二次迁移（应该是幂等的）
        manager.migrate(&conn).unwrap();
        let version_after_second = manager.get_current_version(&conn).unwrap();

        assert_eq!(version_after_first, version_after_second, "重复迁移应该是幂等的");
    }

    #[test]
    fn test_specific_migration_features() {
        let conn = create_test_connection();
        let manager = MigrationManager::new();

        // 执行迁移
        manager.migrate(&conn).unwrap();

        // 测试特定的迁移功能

        // 1. 检查materials表是否有model_id字段（迁移2）
        let has_model_id = conn.prepare("SELECT model_id FROM materials LIMIT 1").is_ok();
        assert!(has_model_id, "materials表应该有model_id字段");

        // 2. 检查template_materials表是否有file_exists字段（迁移3）
        let has_file_exists = conn.prepare("SELECT file_exists FROM template_materials LIMIT 1").is_ok();
        assert!(has_file_exists, "template_materials表应该有file_exists字段");

        // 3. 检查materials表是否有thumbnail_path字段（迁移8）
        let has_materials_thumbnail_path = conn.prepare("SELECT thumbnail_path FROM materials LIMIT 1").is_ok();
        assert!(has_materials_thumbnail_path, "materials表应该有thumbnail_path字段");

        // 4. 检查material_segments表是否有usage_count字段（迁移8）
        let has_usage_count = conn.prepare("SELECT usage_count FROM material_segments LIMIT 1").is_ok();
        assert!(has_usage_count, "material_segments表应该有usage_count字段");

        // 5. 检查template_matching_results表是否有export_count字段（迁移8）
        let has_export_count = conn.prepare("SELECT export_count FROM template_matching_results LIMIT 1").is_ok();
        assert!(has_export_count, "template_matching_results表应该有export_count字段");

        // 6. 验证初始化时就存在的字段
        let has_matching_rule = conn.prepare("SELECT matching_rule FROM track_segments LIMIT 1").is_ok();
        assert!(has_matching_rule, "track_segments表应该有matching_rule字段（初始化时创建）");

        let has_segments_thumbnail_path = conn.prepare("SELECT thumbnail_path FROM material_segments LIMIT 1").is_ok();
        assert!(has_segments_thumbnail_path, "material_segments表应该有thumbnail_path字段（初始化时创建）");
    }
}