Tolan/backend/migrations/migrate_backfill_default_chatbot_scene.py

"""
迁移脚本：给老用户补齐默认 /chatbot 场景（幂等）

默认规则：
- 幕布：fresh_green_wallpaper
- slot_04：stump_stool
- slot_03：french_window

并确保用户拥有：
- fresh_green_wallpaper
- stump_stool
- french_window
"""

import asyncio
import sys
from pathlib import Path
from typing import Any

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

from shop_item_image_utils import scene_image_name_from_shop_row


REQUIRED_TABLES = [
    "users",
    "shop_items",
    "user_assets",
    "user_chatbot_backdrops",
    "user_chatbot_slot_placements",
]

ITEM_CODES = {
    "backdrop": "fresh_green_wallpaper",
    "seat": "stump_stool",
    "window": "french_window",
}


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_exists(conn: aiomysql.Connection, table: str, col: str) -> bool:
    async with conn.cursor() as cursor:
        await cursor.execute(
            """
            SELECT COUNT(*) FROM information_schema.columns
            WHERE table_schema=%s AND table_name=%s AND column_name=%s
            """,
            (settings.MYSQL_DATABASE, table, col),
        )
        row: Any = await cursor.fetchone()
        return bool(row and row[0] > 0)


async def _col_is_bigint(conn: aiomysql.Connection, table: str, col: str) -> bool:
    async with conn.cursor() as cursor:
        await cursor.execute(
            """
            SELECT DATA_TYPE FROM information_schema.columns
            WHERE table_schema=%s AND table_name=%s AND column_name=%s
            """,
            (settings.MYSQL_DATABASE, table, col),
        )
        row: Any = await cursor.fetchone()
        return bool(row and "bigint" in str(row[0]).lower())


async def migrate() -> bool:
    conn: aiomysql.Connection | None = None
    try:
        conn = 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",
            autocommit=False,
        )

        missing = [t for t in REQUIRED_TABLES if not await _table_exists(conn, t)]
        if missing:
            print(f"❌ 缺少必要表: {missing}")
            return False

        ua_bigint = await _col_is_bigint(conn, "user_assets", "user_id")
        ub_bigint = await _col_is_bigint(conn, "user_chatbot_backdrops", "user_id")
        up_bigint = await _col_is_bigint(conn, "user_chatbot_slot_placements", "user_id")

        ua_uid_expr = "u.id" if ua_bigint else "u.user_id"
        ub_uid_expr = "u.id" if ub_bigint else "u.user_id"
        up_uid_expr = "u.id" if up_bigint else "u.user_id"

        has_render_col = await _column_exists(conn, "shop_items", "render_image_name")

        async with conn.cursor() as cursor:
            if has_render_col:
                await cursor.execute(
                    """
                    SELECT item_code, id, image_name, render_image_name, default_z_index
                    FROM shop_items
                    WHERE item_code IN (%s,%s,%s)
                    """,
                    (ITEM_CODES["backdrop"], ITEM_CODES["seat"], ITEM_CODES["window"]),
                )
                rows = await cursor.fetchall()
                by_code = {
                    r[0]: {
                        "id": int(r[1]),
                        "image_name": str(r[2]),
                        "render_image_name": r[3],
                        "z": int(r[4] or 0),
                    }
                    for r in rows
                }
            else:
                await cursor.execute(
                    """
                    SELECT item_code, id, image_name, default_z_index
                    FROM shop_items
                    WHERE item_code IN (%s,%s,%s)
                    """,
                    (ITEM_CODES["backdrop"], ITEM_CODES["seat"], ITEM_CODES["window"]),
                )
                rows = await cursor.fetchall()
                by_code = {
                    r[0]: {
                        "id": int(r[1]),
                        "image_name": str(r[2]),
                        "render_image_name": None,
                        "z": int(r[3] or 0),
                    }
                    for r in rows
                }
            for k, code in ITEM_CODES.items():
                if code not in by_code:
                    print(f"❌ shop_items 缺少 {code}")
                    return False

            seat = by_code[ITEM_CODES["seat"]]
            window = by_code[ITEM_CODES["window"]]
            backdrop = by_code[ITEM_CODES["backdrop"]]

            # 1) 补 user_assets（拥有默认三件，used_qty 先不抬高）
            for role in ("seat", "window", "backdrop"):
                item_id = by_code[ITEM_CODES[role]]["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 {ua_uid_expr}, %s, 1, 0, NOW(), NOW(), NOW()
                    FROM users u
                    WHERE NOT EXISTS (
                      SELECT 1 FROM user_assets ua
                      WHERE ua.user_id = {ua_uid_expr} AND ua.item_id = %s
                    )
                    """,
                    (item_id, item_id),
                )
            print("✅ user_assets 默认三件补齐完成")

            # 2) 补默认幕布
            await cursor.execute(
                f"""
                INSERT INTO user_chatbot_backdrops
                  (user_id, backdrop_item_id, z_index, created_at, updated_at)
                SELECT {ub_uid_expr}, %s, 0, NOW(), NOW()
                FROM users u
                WHERE NOT EXISTS (
                  SELECT 1 FROM user_chatbot_backdrops ub
                  WHERE ub.user_id = {ub_uid_expr}
                )
                """,
                (backdrop["id"],),
            )
            print("✅ user_chatbot_backdrops 默认幕布补齐完成")

            # 3) 槽位摆放：slot_04 放 seat；slot_03 放 window
            # 先通过 user_assets 取每个用户对应 user_asset_id
            await cursor.execute(
                f"""
                INSERT INTO user_chatbot_slot_placements
                  (user_id, backdrop_item_id, slot_code, user_asset_id, item_id, z_index, image_name_snapshot, placed_at, updated_at)
                SELECT
                  {up_uid_expr},
                  %s,
                  'slot_04_row4_col2',
                  ua.id,
                  %s,
                  %s,
                  %s,
                  NOW(),
                  NOW()
                FROM users u
                JOIN user_assets ua
                  ON ua.user_id = {ua_uid_expr}
                 AND ua.item_id = %s
                ON DUPLICATE KEY UPDATE
                  user_asset_id=VALUES(user_asset_id),
                  item_id=VALUES(item_id),
                  z_index=VALUES(z_index),
                  image_name_snapshot=VALUES(image_name_snapshot),
                  updated_at=NOW()
                """,
                (
                    backdrop["id"],
                    seat["id"],
                    seat["z"],
                    scene_image_name_from_shop_row(seat),
                    seat["id"],
                ),
            )

            await cursor.execute(
                f"""
                INSERT INTO user_chatbot_slot_placements
                  (user_id, backdrop_item_id, slot_code, user_asset_id, item_id, z_index, image_name_snapshot, placed_at, updated_at)
                SELECT
                  {up_uid_expr},
                  %s,
                  'slot_03_center_2x3',
                  ua.id,
                  %s,
                  %s,
                  %s,
                  NOW(),
                  NOW()
                FROM users u
                JOIN user_assets ua
                  ON ua.user_id = {ua_uid_expr}
                 AND ua.item_id = %s
                ON DUPLICATE KEY UPDATE
                  user_asset_id=VALUES(user_asset_id),
                  item_id=VALUES(item_id),
                  z_index=VALUES(z_index),
                  image_name_snapshot=VALUES(image_name_snapshot),
                  updated_at=NOW()
                """,
                (
                    backdrop["id"],
                    window["id"],
                    window["z"],
                    scene_image_name_from_shop_row(window),
                    window["id"],
                ),
            )
            print("✅ user_chatbot_slot_placements 默认槽位补齐完成")

            # 4) used_qty 与默认摆放同步（seat/window 标记已使用1个）
            await cursor.execute(
                f"""
                UPDATE user_assets ua
                JOIN users u ON ua.user_id = {ua_uid_expr}
                SET ua.used_qty = GREATEST(ua.used_qty, 1), ua.updated_at = NOW()
                WHERE ua.item_id IN (%s,%s)
                """,
                (seat["id"], window["id"]),
            )

        await conn.commit()
        print("🎉 migrate_backfill_default_chatbot_scene 完成")
        return True
    except Exception as exc:  # noqa: BLE001
        print(f"❌ 迁移失败: {exc}")
        if conn is not None:
            try:
                await conn.rollback()
            except Exception:
                pass
        return False
    finally:
        if conn is not None:
            conn.close()


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续为老用户回填默认/chatbot场景? (y/n): ").strip().lower()
    if confirm != "y":
        sys.exit(0)
    ok = asyncio.run(migrate())
    sys.exit(0 if ok else 1)