#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
模块名称：migrate_sqlite_to_mysql.py
模块描述：SQLite → MySQL 数据迁移脚本
    将本地 SQLite 数据库（backend/data/test_platform.db）的数据迁移到远程 MySQL。
    支持断点续传：已存在的记录自动跳过（ON DUPLICATE KEY UPDATE）。

作者：czj
创建日期：2026-07-28
"""

import asyncio
import json
import logging
import os
import sys
from datetime import datetime

from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# ==================== 数据库连接配置 ====================

# SQLite（本地数据源）
SQLITE_PATH = os.getenv("SQLITE_PATH", "sqlite+aiosqlite:///./data/test_platform.db")

# MySQL（目标数据库）
MYSQL_HOST = os.getenv("MYSQL_HOST", "192.168.5.60")
MYSQL_PORT = os.getenv("MYSQL_PORT", "3307")
MYSQL_USER = os.getenv("MYSQL_USER", "platapp")
MYSQL_PASSWORD = os.getenv("MYSQL_PASSWORD", "PlatApp2026")
MYSQL_DATABASE = os.getenv("MYSQL_DATABASE", "plat_auto_test")
MYSQL_URL = (
    f"mysql+aiomysql://{MYSQL_USER}:{MYSQL_PASSWORD}"
    f"@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}?charset=utf8mb4"
)

# ==================== 表迁移顺序（按外键依赖排序） ====================

TABLES = [
    # (表名, 主键字段, 是否需要处理 JSON)
    ("modules", "id", False),
    ("test_cases", "id", True),       # 含 JSON 字段（steps, selectors 等）
    ("executions", "id", True),       # 含 JSON 字段
    ("case_results", "id", False),    # 依赖 execution.id
    ("case_dependencies", "id", False),  # 依赖 test_case.id
    ("security_configs", "id", False),
    ("vulnerability_results", "id", False),  # 依赖 security_configs.id
]


async def get_sqlite_data(sqlite_engine, table: str) -> list:
    """从 SQLite 读取表数据"""
    async with sqlite_engine.connect() as conn:
        result = await conn.execute(text(f"SELECT * FROM {table}"))
        columns = result.keys()
        rows = []
        for row in result.fetchall():
            rows.append(dict(zip(columns, row)))
        logger.info(f"  SQLite [{table}]: {len(rows)} 行")
        return rows


def serialize_value(value, is_json: bool = False):
    """序列化值，处理特殊类型"""
    if value is None:
        return None
    if is_json and not isinstance(value, (str, int, float, bool)):
        # JSON 字段：如果是 dict/list 则序列化，否则保持原样
        if isinstance(value, (dict, list)):
            return json.dumps(value, ensure_ascii=False)
    # datetime 类型转为字符串
    if isinstance(value, datetime):
        return value.strftime("%Y-%m-%d %H:%M:%S")
    # bytes 类型转为字符串
    if isinstance(value, bytes):
        return value.decode("utf-8")
    return value


async def migrate_table(mysql_engine, table: str, pk_field: str, rows: list, is_json_table: bool) -> int:
    """将数据写入 MySQL"""
    if not rows:
        return 0

    # 获取列名（从第一行数据）
    columns = list(rows[0].keys())

    # 构建 INSERT ... ON DUPLICATE KEY UPDATE 语句
    col_list = ", ".join([f"`{c}`" for c in columns])
    placeholders = ", ".join([f":{c}" for c in columns])
    update_parts = ", ".join([f"`{c}`=VALUES(`{c}`)" for c in columns if c != pk_field])

    sql = f"""
        INSERT INTO `{table}` ({col_list})
        VALUES ({placeholders})
        ON DUPLICATE KEY UPDATE {update_parts}
    """

    # 序列化数据
    serialized_rows = []
    for row in rows:
        serialized = {}
        for col, val in row.items():
            # 判断是否是 JSON 字段 - 通过表名 + 列名判断
            is_json_col = is_json_table and col in ("steps", "selectors", "params", "result", "config", "tags", "metadata", "extra")
            serialized[col] = serialize_value(val, is_json_col)
        serialized_rows.append(serialized)

    async with mysql_engine.connect() as conn:
        for row in serialized_rows:
            await conn.execute(text(sql), row)
        await conn.commit()

    logger.info(f"  MySQL [{table}]: {len(rows)} 行写入成功")
    return len(rows)


async def verify_migration(mysql_engine):
    """验证迁移后的数据"""
    logger.info("")
    logger.info("=" * 50)
    logger.info("验证迁移结果")
    logger.info("=" * 50)

    async with mysql_engine.connect() as conn:
        for table, pk_field, _ in TABLES:
            result = await conn.execute(text(f"SELECT COUNT(*) FROM `{table}`"))
            count = result.scalar()
            logger.info(f"  {table}: {count} 行")

        # 验证用例完整性
        result = await conn.execute(text("SELECT COUNT(*) FROM test_cases"))
        case_count = result.scalar()

        result = await conn.execute(text("SELECT COUNT(*) FROM modules"))
        module_count = result.scalar()

    logger.info("")
    logger.info(f"✅ 迁移完成: {case_count} 个用例, {module_count} 个模块")


async def main():
    """主迁移流程"""
    start_time = datetime.now()
    logger.info("=" * 50)
    logger.info("SQLite → MySQL 数据迁移")
    logger.info(f"来源: {SQLITE_PATH}")
    logger.info(f"目标: {MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}")
    logger.info("=" * 50)

    # 创建引擎
    sqlite_engine = create_async_engine(SQLITE_PATH)
    mysql_engine = create_async_engine(MYSQL_URL)

    # 迁移前：关闭外键检查（避免自引用外键如 modules.parent_id 的顺序问题）
    async with mysql_engine.connect() as conn:
        await conn.execute(text("SET FOREIGN_KEY_CHECKS = 0"))
        await conn.commit()
    logger.info("已关闭 MySQL 外键检查（迁移期间）")

    total_rows = 0
    try:
        # 逐表迁移
        for table, pk_field, is_json in TABLES:
            logger.info("")
            logger.info(f"📦 迁移表: {table}")

            # 1. 从 SQLite 读取
            rows = await get_sqlite_data(sqlite_engine, table)

            # 2. 写入 MySQL
            count = await migrate_table(mysql_engine, table, pk_field, rows, is_json)
            total_rows += count

        # 重新开启外键检查
        async with mysql_engine.connect() as conn:
            await conn.execute(text("SET FOREIGN_KEY_CHECKS = 1"))
            await conn.commit()
        logger.info("已重新开启 MySQL 外键检查")

        # 3. 验证
        await verify_migration(mysql_engine)

    except Exception as e:
        logger.error(f"❌ 迁移失败: {str(e)}")
        raise
    finally:
        await sqlite_engine.dispose()
        await mysql_engine.dispose()

    elapsed = (datetime.now() - start_time).total_seconds()
    logger.info("")
    logger.info(f"⏱ 总耗时: {elapsed:.1f} 秒")
    logger.info(f"📊 总行数: {total_rows}")

    return total_rows


if __name__ == "__main__":
    asyncio.run(main())