Tolan/backend/migrations/migrate_add_avatars.py

"""
数据库迁移脚本：创建形象相关表，并为 user_ui_states 添加 current_avatar_id
运行: python migrate_add_avatars.py
"""
import asyncio
import sys
from typing import Any

import aiomysql

from config import settings


DEFAULT_AVATAR_ID = "avatar_default"
REGISTERED_AVATAR_ID = "avatar_registered_cap"

# 默认形象 GIF 资源
DEFAULT_AVATAR_ASSETS = [
    (DEFAULT_AVATAR_ID, "default", "https://www.lyrria.com/img/01-Default.gif"),
    (DEFAULT_AVATAR_ID, "positive", "https://www.lyrria.com/img/04-Positive-emotion.gif"),
    (DEFAULT_AVATAR_ID, "negative", "https://www.lyrria.com/img/05-Negative-emotion.gif"),
    (DEFAULT_AVATAR_ID, "talking", "https://www.lyrria.com/img/06-talking.gif"),
]

# 注册赠送形象 GIF 资源
REGISTERED_AVATAR_ASSETS = [
    (
        REGISTERED_AVATAR_ID,
        "default",
        "https://www.lyrria.com/img/01-Default-Registered.gif",
    ),
    (
        REGISTERED_AVATAR_ID,
        "positive",
        "https://www.lyrria.com/img/04-Positive-emotion-Registered.gif",
    ),
    (
        REGISTERED_AVATAR_ID,
        "negative",
        "https://www.lyrria.com/img/05-Negative-emotion-Registered.gif",
    ),
    (
        REGISTERED_AVATAR_ID,
        "talking",
        "https://www.lyrria.com/img/06-talking-Registered.gif",
    ),
]

# 统一的字符序（用于跨表 join 避免字符集/排序规则不一致报错）
COLLATE_NAME = "utf8mb4_unicode_ci"


async def migrate_create_avatar_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. avatars 表
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.tables
                WHERE table_schema = %s AND table_name = 'avatars'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row: Any = await cursor.fetchone()
            avatars_exists = bool(exists_row and exists_row[0] > 0)

            if not avatars_exists:
                print("🔄 创建 avatars 表...")
                await cursor.execute(
                    """
                    CREATE TABLE avatars (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        avatar_id VARCHAR(64) NOT NULL UNIQUE COMMENT '形象唯一标识',
                        avatar_name VARCHAR(100) NOT NULL COMMENT '形象名称',
                        description VARCHAR(255) NULL COMMENT '形象描述',
                        is_active TINYINT(1) NOT NULL DEFAULT 1 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_avatar_is_active (is_active),
                        INDEX idx_avatar_id (avatar_id)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='形象配置表';
                    """
                )
                print("✅ avatars 表创建完成")
            else:
                print("✅ avatars 表已存在")

            # 2. avatar_assets 表
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.tables
                WHERE table_schema = %s AND table_name = 'avatar_assets'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row = await cursor.fetchone()
            avatar_assets_exists = bool(exists_row and exists_row[0] > 0)

            if not avatar_assets_exists:
                print("🔄 创建 avatar_assets 表...")
                await cursor.execute(
                    """
                    CREATE TABLE avatar_assets (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        avatar_id VARCHAR(64) NOT NULL COMMENT '形象ID',
                        state ENUM('default', 'positive', 'negative', 'talking', 'idle', 'sleep') NOT NULL DEFAULT 'default' COMMENT '动作/情绪状态',
                        image_url VARCHAR(500) NOT NULL COMMENT 'GIF 资源 URL',
                        thumbnail_url VARCHAR(500) NULL COMMENT '缩略图URL',
                        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 '更新时间',
                        UNIQUE KEY uk_avatar_asset_state (avatar_id, state),
                        INDEX idx_avatar_asset_avatar_id (avatar_id)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='形象动作资源表';
                    """
                )
                print("✅ avatar_assets 表创建完成")
            else:
                print("✅ avatar_assets 表已存在")

            # 3. user_avatars 表
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.tables
                WHERE table_schema = %s AND table_name = 'user_avatars'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row = await cursor.fetchone()
            user_avatars_exists = bool(exists_row and exists_row[0] > 0)

            if not user_avatars_exists:
                print("🔄 创建 user_avatars 表...")
                await cursor.execute(
                    """
                    CREATE TABLE user_avatars (
                        id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '主键',
                        user_id VARCHAR(64) NOT NULL COMMENT '用户ID',
                        avatar_id VARCHAR(64) NOT NULL COMMENT '形象ID',
                        obtained_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '获得时间',
                        obtained_type ENUM('default', 'register', 'level_up', 'mission', 'purchase', 'gift', 'event') 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_avatar (user_id, avatar_id),
                        INDEX idx_user_avatar_user_obtained (user_id, obtained_at),
                        INDEX idx_user_avatar_avatar_id (avatar_id)
                    ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='用户已拥有的形象';
                    """
                )
                print("✅ user_avatars 表创建完成")
            else:
                print("✅ user_avatars 表已存在")

            # 4. user_ui_states 表添加 current_avatar_id
            await cursor.execute(
                """
                SELECT COUNT(*) FROM information_schema.columns
                WHERE table_schema = %s
                  AND table_name = 'user_ui_states'
                  AND column_name = 'current_avatar_id'
                """,
                (settings.MYSQL_DATABASE,),
            )
            exists_row = await cursor.fetchone()
            column_exists = bool(exists_row and exists_row[0] > 0)

            if not column_exists:
                print("🔄 为 user_ui_states 添加 current_avatar_id 字段...")
                await cursor.execute(
                    """
                    ALTER TABLE user_ui_states
                    ADD COLUMN current_avatar_id VARCHAR(64) NULL COMMENT '当前使用的形象ID(关联 avatars.avatar_id)'
                    AFTER has_shared_avatar
                    """
                )
                print("✅ current_avatar_id 字段添加完成")
            else:
                print("✅ current_avatar_id 字段已存在")

            await conn.commit()

            # 5. 初始化 avatars 数据
            print("🔄 初始化 avatars 基础数据...")
            await cursor.execute(
                """
                INSERT IGNORE INTO avatars (avatar_id, avatar_name, description, is_active)
                VALUES 
                    (%s, %s, %s, 1),
                    (%s, %s, %s, 1)
                """,
                (
                    DEFAULT_AVATAR_ID,
                    "默认形象",
                    "基础形象（未注册/通用）",
                    REGISTERED_AVATAR_ID,
                    "注册赠送形象",
                    "注册奖励形象（带帽子）",
                ),
            )
            await conn.commit()
            print("✅ avatars 基础数据初始化完成")

            # 6. 初始化 avatar_assets 数据
            print("🔄 初始化 avatar_assets 数据...")
            asset_rows = DEFAULT_AVATAR_ASSETS + REGISTERED_AVATAR_ASSETS
            for avatar_id, state, image_url in asset_rows:
                await cursor.execute(
                    """
                    INSERT IGNORE INTO avatar_assets (avatar_id, state, image_url)
                    VALUES (%s, %s, %s)
                    """,
                    (avatar_id, state, image_url),
                )
            await conn.commit()
            print(f"✅ 已初始化 {len(asset_rows)} 条 avatar_assets 数据")

            # 7. 为所有用户分配默认形象
            print("🔄 为用户添加默认形象拥有记录...")
            await cursor.execute(
                """
                INSERT IGNORE INTO user_avatars (user_id, avatar_id, obtained_type, obtained_source)
                SELECT user_id, %s, 'default', 'system_init'
                FROM users
                """,
                (DEFAULT_AVATAR_ID,),
            )
            await conn.commit()
            print(f"✅ 默认形象写入 user_avatars，受影响行数: {cursor.rowcount}")

            # 8. 为注册用户分配注册赠送形象
            print("🔄 为注册用户添加注册赠送形象...")
            await cursor.execute(
                """
                INSERT IGNORE INTO user_avatars (user_id, avatar_id, obtained_type, obtained_source)
                SELECT user_id, %s, 'register', 'register_bonus'
                FROM users
                WHERE user_type = 'registered'
                """,
                (REGISTERED_AVATAR_ID,),
            )
            await conn.commit()
            print(f"✅ 注册赠送形象写入 user_avatars，受影响行数: {cursor.rowcount}")

            # 9. 设置当前形象（注册用户优先注册赠送，其余默认）
            print("🔄 设置 user_ui_states.current_avatar_id ...")
            # 注意：部分历史表的字符集/排序规则不一致，使用显式 COLLATE 避免 1267 错误
            await cursor.execute(
                f"""
                UPDATE user_ui_states ui
                JOIN users u 
                  ON ui.user_id COLLATE {COLLATE_NAME} = u.user_id COLLATE {COLLATE_NAME}
                SET ui.current_avatar_id = %s
                WHERE ui.current_avatar_id IS NULL
                  AND u.user_type = 'registered'
                """,
                (REGISTERED_AVATAR_ID,),
            )
            await conn.commit()

            await cursor.execute(
                f"""
                UPDATE user_ui_states
                SET current_avatar_id = %s
                WHERE current_avatar_id IS NULL
                """,
                (DEFAULT_AVATAR_ID,),
            )
            await conn.commit()
            print("✅ current_avatar_id 初始化完成")

        conn.close()
        return True
    except Exception:
        import traceback

        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_avatar_tables())
    if not ok:
        sys.exit(1)