Tolan/backend/migrations/migrate_add_background_images.py

"""
数据库迁移脚本：创建背景图片相关表并初始化数据
运行: python migrate_add_background_images.py
"""
import asyncio
import sys
from typing import Any

import aiomysql

from config import settings


async def migrate_create_background_images_tables() -> bool:
    """创建背景图片相关表并初始化数据"""
    try:
        conn: aiomysql.Connection = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset="utf8mb4",
        )

        async with conn.cursor() as cursor:
            # 1. 检查并创建 background_images 表
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.tables
                WHERE table_schema = %s AND table_name = 'background_images'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row: Any = await cursor.fetchone()
            bg_images_exists = bool(exists_row and exists_row[0] > 0)

            if not bg_images_exists:
                print("🔄 创建 background_images 表...")
                await cursor.execute(
                    """
                    CREATE TABLE background_images (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        image_id VARCHAR(64) NOT NULL UNIQUE COMMENT '背景图片唯一标识',
                        image_name VARCHAR(100) NOT NULL COMMENT '背景图片名称',
                        image_url VARCHAR(500) NOT NULL COMMENT '背景图片URL',
                        thumbnail_url VARCHAR(500) NULL COMMENT '缩略图URL',
                        unlock_type ENUM('default', 'level', 'purchase', 'mission') NOT NULL DEFAULT 'default' COMMENT '解锁类型',
                        unlock_level INT NULL COMMENT '解锁所需等级（unlock_type=level时）',
                        unlock_mission_type VARCHAR(64) NULL COMMENT '解锁所需任务类型（unlock_type=mission时）',
                        display_order INT NOT NULL DEFAULT 0 COMMENT '显示顺序',
                        is_active TINYINT(1) NOT NULL DEFAULT 1 COMMENT '是否启用',
                        description VARCHAR(255) NULL COMMENT '背景图片描述',
                        metadata JSON NULL COMMENT '扩展配置',
                        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                        updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                        INDEX idx_image_id (image_id),
                        INDEX idx_unlock_type (unlock_type),
                        INDEX idx_display_order (display_order)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='背景图片配置表';
                    """
                )
                print("✅ background_images 表创建完成")
            else:
                print("✅ background_images 表已存在")

            # 2. 检查并创建 user_background_images 表
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.tables
                WHERE table_schema = %s AND table_name = 'user_background_images'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row = await cursor.fetchone()
            user_bg_images_exists = bool(exists_row and exists_row[0] > 0)

            if not user_bg_images_exists:
                print("🔄 创建 user_background_images 表...")
                await cursor.execute(
                    """
                    CREATE TABLE user_background_images (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
                        image_id VARCHAR(64) NOT NULL COMMENT '背景图片ID',
                        obtained_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '获得时间',
                        obtained_type ENUM('default', 'level_up', 'purchase', 'mission', 'gift') NOT NULL DEFAULT 'default' COMMENT '获得方式',
                        obtained_source VARCHAR(64) NULL COMMENT '获得来源（任务类型、等级等）',
                        metadata JSON NULL COMMENT '扩展信息',
                        created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                        UNIQUE KEY uk_user_background_image (user_id, image_id),
                        INDEX idx_user_id (user_id),
                        INDEX idx_image_id (image_id),
                        INDEX idx_user_obtained (user_id, obtained_at)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='用户背景图片关联表';
                    """
                )
                print("✅ user_background_images 表创建完成")
            else:
                print("✅ user_background_images 表已存在")

            # 3. 检查并添加 users 表的 current_background_image_id 字段
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.columns
                WHERE table_schema = %s 
                  AND table_name = 'users'
                  AND column_name = 'current_background_image_id'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row = await cursor.fetchone()
            column_exists = bool(exists_row and exists_row[0] > 0)

            if not column_exists:
                print("🔄 为 users 表添加 current_background_image_id 字段...")
                await cursor.execute(
                    """
                    ALTER TABLE users
                    ADD COLUMN current_background_image_id VARCHAR(64) NULL COMMENT '当前使用的背景图片ID'
                    AFTER last_active_at
                    """
                )
                print("✅ current_background_image_id 字段添加完成")
            else:
                print("✅ current_background_image_id 字段已存在")

            await conn.commit()

            # 4. 初始化背景图片数据
            print("🔄 初始化背景图片数据...")
            
            # 检查默认背景图片是否存在
            await cursor.execute(
                "SELECT COUNT(*) FROM background_images WHERE image_id = 'default'"
            )
            default_exists = (await cursor.fetchone())[0] > 0
            
            if not default_exists:
                await cursor.execute(
                    """
                    INSERT INTO background_images 
                    (image_id, image_name, image_url, unlock_type, display_order, description)
                    VALUES 
                    ('default', '默认背景', 'https://assets.lyrria.com/init.jpg', 'default', 1, '默认背景图片')
                    """
                )
                print("✅ 默认背景图片已添加")
            else:
                print("✅ 默认背景图片已存在")

            # 检查新年背景图片是否存在
            await cursor.execute(
                "SELECT COUNT(*) FROM background_images WHERE image_id = 'happynewyear'"
            )
            newyear_exists = (await cursor.fetchone())[0] > 0
            
            if not newyear_exists:
                await cursor.execute(
                    """
                    INSERT INTO background_images 
                    (image_id, image_name, image_url, unlock_type, display_order, description)
                    VALUES 
                    ('happynewyear', '新年背景', 'https://assets.lyrria.com/bg-happynewyear.png', 'default', 2, '新年背景图片')
                    """
                )
                print("✅ 新年背景图片已添加")
            else:
                print("✅ 新年背景图片已存在")

            await conn.commit()

            # 5. 为所有现有用户设置默认背景图片
            print("🔄 为现有用户设置默认背景图片...")
            
            # 更新所有用户的 current_background_image_id 为 'default'（如果为 NULL）
            await cursor.execute(
                """
                UPDATE users 
                SET current_background_image_id = 'default'
                WHERE current_background_image_id IS NULL
                """
            )
            updated_users = cursor.rowcount
            print(f"✅ 已为 {updated_users} 个用户设置默认背景图片")

            # 6. 为所有现有用户添加默认背景图片到已获得列表
            print("🔄 为现有用户添加默认背景图片到已获得列表...")
            
            await cursor.execute(
                """
                INSERT IGNORE INTO user_background_images 
                (user_id, image_id, obtained_type, obtained_source)
                SELECT user_id, 'default', 'default', 'system_init'
                FROM users
                """
            )
            added_records = cursor.rowcount
            print(f"✅ 已为 {added_records} 个用户添加默认背景图片到已获得列表")

            await conn.commit()
            print("✅ 迁移完成")

        await conn.ensure_closed()
        return True
    except Exception as exc:  # noqa: BLE001
        import traceback

        print(f"❌ 迁移失败: {exc}")
        traceback.print_exc()
        return False


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续创建背景图片相关表并初始化数据? (y/n): ")

    if confirm.lower() != "y":
        sys.exit(0)

    ok = asyncio.run(migrate_create_background_images_tables())
    if not ok:
        sys.exit(1)