#!/usr/bin/env python3
"""将本地 MySQL 数据迁移到远程 MySQL。

默认配置：
  源库：127.0.0.1:3306/weread_dashboard
  目标库：124.222.38.85:13306/weread_dashboard

密码优先从环境变量读取：
  SOURCE_DB_PASSWORD / TARGET_DB_PASSWORD

脚本不会创建数据库或表；请先在目标 MySQL 执行 schema_mysql.sql。
目标表会被清空后重新写入，因此执行前请确认目标库中没有需要保留的数据。
"""

import argparse
import os
import sys
from typing import Dict, Iterable, List, Tuple

import pymysql


def db_config(prefix: str, defaults: Dict[str, object]) -> Dict[str, object]:
    return {
        "host": os.getenv(f"{prefix}_DB_HOST", defaults["host"]),
        "port": int(os.getenv(f"{prefix}_DB_PORT", defaults["port"])),
        "user": os.getenv(f"{prefix}_DB_USER", defaults["user"]),
        "password": os.getenv(f"{prefix}_DB_PASSWORD", ""),
        "database": os.getenv(f"{prefix}_DB_NAME", defaults["database"]),
        "charset": "utf8mb4",
        "autocommit": False,
    }


SOURCE = db_config("SOURCE", {
    "host": "127.0.0.1", "port": 3306, "user": "root", "database": "weread_dashboard"
})
TARGET = db_config("TARGET", {
    "host": "124.222.38.85", "port": 13306, "user": "root", "database": "weread_dashboard"
})


def q(identifier: str) -> str:
    """安全引用 MySQL 标识符。表名和列名均来自数据库元数据。"""
    return "`" + identifier.replace("`", "``") + "`"


def tables(conn) -> List[str]:
    with conn.cursor() as cur:
        cur.execute("SHOW FULL TABLES WHERE Table_type = 'BASE TABLE'")
        return [row[0] for row in cur.fetchall()]


def columns(conn, table: str) -> List[str]:
    with conn.cursor() as cur:
        cur.execute(f"SHOW COLUMNS FROM {q(table)}")
        return [row[0] for row in cur.fetchall()]


def row_count(conn, table: str) -> int:
    with conn.cursor() as cur:
        cur.execute(f"SELECT COUNT(*) FROM {q(table)}")
        return int(cur.fetchone()[0])


def migrate_table(source, target, table: str, batch_size: int) -> Tuple[int, int]:
    source_cols = columns(source, table)
    target_cols = set(columns(target, table))
    common = [col for col in source_cols if col in target_cols]
    if not common:
        raise RuntimeError(f"表 {table} 没有可写入的同名列")

    col_sql = ", ".join(q(col) for col in common)
    placeholders = ", ".join(["%s"] * len(common))
    insert_sql = f"INSERT INTO {q(table)} ({col_sql}) VALUES ({placeholders})"

    copied = 0
    with source.cursor() as sc, target.cursor() as tc:
        sc.execute(f"SELECT {col_sql} FROM {q(table)}")
        while True:
            batch = sc.fetchmany(batch_size)
            if not batch:
                break
            tc.executemany(insert_sql, batch)
            copied += len(batch)
    return copied, len(source_cols) - len(common)


def migrate(batch_size: int, only: Iterable[str] = ()) -> None:
    source = pymysql.connect(**SOURCE)
    target = pymysql.connect(**TARGET)
    try:
        source_tables = tables(source)
        target_tables = set(tables(target))
        selected = list(only) if only else source_tables
        missing = [t for t in selected if t not in target_tables]
        if missing:
            raise RuntimeError("目标库缺少表：" + ", ".join(missing))

        print(f"源库: {SOURCE['host']}:{SOURCE['port']}/{SOURCE['database']}")
        print(f"目标库: {TARGET['host']}:{TARGET['port']}/{TARGET['database']}")
        print(f"迁移表数: {len(selected)}，批大小: {batch_size}\n")

        with target.cursor() as cur:
            cur.execute("SET FOREIGN_KEY_CHECKS = 0")
            for table in selected:
                before = row_count(source, table)
                cur.execute(f"TRUNCATE TABLE {q(table)}")
                copied, skipped_cols = migrate_table(source, target, table, batch_size)
                target.commit()
                after = row_count(target, table)
                note = f"，跳过源表独有列 {skipped_cols} 个" if skipped_cols else ""
                status = "✓" if before == after == copied else "!"
                print(f"{status} {table:32s} 源 {before:8d} / 目标 {after:8d}{note}")
            cur.execute("SET FOREIGN_KEY_CHECKS = 1")
            target.commit()
    except Exception:
        target.rollback()
        raise
    finally:
        source.close()
        target.close()


def main() -> int:
    parser = argparse.ArgumentParser(description="本地 MySQL → 远程 MySQL 数据迁移")
    parser.add_argument("--batch-size", type=int, default=1000)
    parser.add_argument("--tables", nargs="*", help="只迁移指定表；不传则迁移源库全部表")
    args = parser.parse_args()
    if args.batch_size < 1:
        parser.error("--batch-size 必须大于 0")
    try:
        migrate(args.batch_size, args.tables or ())
        print("\n迁移完成。")
        return 0
    except pymysql.MySQLError as exc:
        print(f"MySQL 错误：{exc}", file=sys.stderr)
        return 2
    except Exception as exc:
        print(f"迁移失败：{exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())