Tolan/backend/migrations/migrate_scenario_gift_items_to_all_users.py

"""
迁移脚本：仅为「已触发过特殊场景」的老用户补齐对应赠送商品

规则：
- scene=window  -> coffee_mug（咖啡杯）
- scene=reading -> books（书）
- scene=eating  -> snacks（零食）

触发依据：
- user_special_scenario_topic_logs.scenario_key
- 同一用户同一场景只补对应那一件商品（已拥有则跳过）

执行：
python migrations/migrate_scenario_gift_items_to_all_users.py
或：
python migrations/migrate_scenario_gift_items_to_all_users.py --force
"""

import asyncio
import sys
from pathlib import Path
from typing import Any, Dict, List

import aiomysql

backend_dir = Path(__file__).resolve().parents[1]
if str(backend_dir) not in sys.path:
    sys.path.insert(0, str(backend_dir))

from config import settings


async def _table_exists(conn: aiomysql.Connection, table_name: str) -> bool:
    async with conn.cursor() as cursor:
        await cursor.execute(
            """
            SELECT COUNT(*) FROM information_schema.tables
            WHERE table_schema = %s AND table_name = %s
            """,
            (settings.MYSQL_DATABASE, table_name),
        )
        row: Any = await cursor.fetchone()
        return bool(row and row[0] > 0)


async def _column_data_type(conn: aiomysql.Connection, table_name: str, column_name: str) -> str:
    async with conn.cursor() as cursor:
        await cursor.execute(
            """
            SELECT DATA_TYPE, COLUMN_TYPE
            FROM information_schema.columns
            WHERE table_schema = %s AND table_name = %s AND column_name = %s
            """,
            (settings.MYSQL_DATABASE, table_name, column_name),
        )
        row: Any = await cursor.fetchone()
        if not row:
            return ""
        data_type = row[0] or ""
        column_type = row[1] or ""
        return f"{data_type}|{column_type}"


async def migrate() -> 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",
        )

        required_tables = ["users", "shop_items", "user_assets", "user_special_scenario_topic_logs"]
        missing = [t for t in required_tables if not await _table_exists(conn, t)]
        if missing:
            print(f"❌ 缺少必要表：{missing}。请先创建对应表后再执行。")
            await conn.ensure_closed()
            return False

        user_id_type_hint = await _column_data_type(conn, "user_assets", "user_id")
        user_id_is_bigint = "bigint" in user_id_type_hint.lower()
        user_id_expr = "u.id" if user_id_is_bigint else "u.user_id"
        if not user_id_type_hint:
            print("❌ 未能读取 user_assets.user_id 字段类型，无法确定映射方式。")
            await conn.ensure_closed()
            return False

        gift_items: List[Dict[str, str]] = [
            {"scenario_key": "window", "item_code": "coffee_mug", "note": "window 窗前"},
            {"scenario_key": "reading", "item_code": "books", "note": "reading 看书"},
            {"scenario_key": "eating", "item_code": "snacks", "note": "eating 吃饭"},
        ]
        item_codes = [it["item_code"] for it in gift_items]

        async with conn.cursor() as cursor:
            await cursor.execute(
                f"""
                SELECT `id`, `item_code`, `image_name`
                FROM `shop_items`
                WHERE `item_code` IN ({", ".join(["%s"] * len(item_codes))})
                """,
                tuple(item_codes),
            )
            rows: List[Any] = list(await cursor.fetchall())
            if not rows:
                print("❌ 在 shop_items 中没有找到情景赠送商品。请先执行 chatbot_shop_items_seed.sql。")
                await conn.ensure_closed()
                return False

            item_code_to_item_id: Dict[str, int] = {r[1]: r[0] for r in rows}
            missing_item_codes = [c for c in item_codes if c not in item_code_to_item_id]
            if missing_item_codes:
                print(f"❌ 缺少物品：item_code={missing_item_codes}")
                await conn.ensure_closed()
                return False

            inserted_total = 0
            for it in gift_items:
                item_id = item_code_to_item_id[it["item_code"]]
                scenario_key = it["scenario_key"]
                print(
                    f"🔄 补齐情景赠送：scenario={scenario_key}, item_code={it['item_code']} ({it['note']}) item_id={item_id}"
                )
                await cursor.execute(
                    f"""
                    INSERT INTO `user_assets`
                      (`user_id`, `item_id`, `owned_qty`, `used_qty`, `last_acquired_at`, `created_at`, `updated_at`)
                    SELECT
                      {user_id_expr},
                      %s,
                      1,
                      0,
                      NOW(),
                      NOW(),
                      NOW()
                    FROM `users` u
                    WHERE EXISTS (
                      SELECT 1
                      FROM `user_special_scenario_topic_logs` slog
                      WHERE slog.`user_id` = u.`user_id`
                        AND slog.`scenario_key` = %s
                    )
                    AND NOT EXISTS (
                      SELECT 1
                      FROM `user_assets` ua
                      WHERE ua.`item_id` = %s
                        AND ua.`user_id` = {user_id_expr}
                    )
                    """,
                    (item_id, scenario_key, item_id),
                )
                inserted = cursor.rowcount or 0
                inserted_total += inserted
                print(f"✅ 完成：{it['item_code']}，新增资产行数={inserted}")

            await conn.commit()
            print(f"🎉 全部完成，总共新增资产行数={inserted_total}")

        await conn.ensure_closed()
        return True
    except Exception as exc:  # noqa: BLE001
        print(f"❌ 迁移失败: {exc}")
        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: Any = asyncio.run(migrate())
    sys.exit(0 if ok else 1)