#!/usr/bin/env python3
"""
将 InfoGenie Go 后端使用的 MySQL 表从「测试库」全量同步到「生产库」。
同步时在目标库对每张表执行：DROP TABLE IF EXISTS → 按源库 SHOW CREATE TABLE 重建 → 插入源库数据，
确保列结构与测试库一致（旧生产表结构会被覆盖）。

使用前务必备份生产库。会删除并重建生产库中同名表。

环境变量（未设置时沿用仓库内 test/test_mysql_connect.py 的约定）：
  源库：DB_SRC_HOST, DB_SRC_PORT, DB_SRC_NAME, DB_SRC_USER, DB_SRC_PASSWORD
  目标：DB_DST_HOST, DB_DST_PORT, DB_DST_NAME, DB_DST_USER, DB_DST_PASSWORD

示例（PowerShell）：
  $env:DB_DST_HOST='192.168.1.100'
  $env:DB_DST_NAME='infogenie'
  $env:DB_DST_USER='infogenie'
  $env:DB_DST_PASSWORD='***'
  python scripts/sync_mysql_test_to_prod.py --dry-run
  python scripts/sync_mysql_test_to_prod.py --execute
"""

from __future__ import annotations

import argparse
import os
import sys
from typing import Any, List, Sequence, Tuple

try:
    import pymysql
except ImportError:
    print("缺少依赖: pip install pymysql", file=sys.stderr)
    sys.exit(2)

# 与 internal/database/mysql.go AutoMigrate 一致
TABLES: Tuple[str, ...] = (
    "ai_configs",
    "site_60s_disabled",
    "site_ai_runtime",
    "site_60s_upstream",
    "site_ai_model_disabled",
    "site_feature_card_clicks",
)

DEFAULT_SRC = {
    "host": "10.1.1.100",
    "port": 3306,
    "name": "infogenie-test",
    "user": "infogenie-test",
    "password": "infogenie-test",
}  # 仅作文档占位；实际以 DB_SRC_* 为准

DEFAULT_DST = {
    "host": "192.168.1.100",
    "port": 3306,
    "name": "infogenie",
    "user": "infogenie",
    "password": "infogenie",
}


def _env_int(key: str, default: int) -> int:
    raw = os.environ.get(key)
    if raw is None or str(raw).strip() == "":
        return default
    return int(str(raw).strip())


def conn_params_src() -> dict:
    return {
        "host": os.environ.get("DB_SRC_HOST", DEFAULT_SRC["host"]),
        "port": _env_int("DB_SRC_PORT", DEFAULT_SRC["port"]),
        "database": os.environ.get("DB_SRC_NAME", DEFAULT_SRC["name"]),
        "user": os.environ.get("DB_SRC_USER", DEFAULT_SRC["user"]),
        "password": os.environ.get("DB_SRC_PASSWORD", DEFAULT_SRC["password"]),
    }


def conn_params_dst() -> dict:
    return {
        "host": os.environ.get("DB_DST_HOST", DEFAULT_DST["host"]),
        "port": _env_int("DB_DST_PORT", DEFAULT_DST["port"]),
        "database": os.environ.get("DB_DST_NAME", DEFAULT_DST["name"]),
        "user": os.environ.get("DB_DST_USER", DEFAULT_DST["user"]),
        "password": os.environ.get("DB_DST_PASSWORD", DEFAULT_DST["password"]),
    }


def open_conn(**kw: Any):
    return pymysql.connect(
        host=kw["host"],
        port=int(kw["port"]),
        user=kw["user"],
        password=kw["password"],
        database=kw["database"],
        charset="utf8mb4",
        connect_timeout=15,
        cursorclass=pymysql.cursors.Cursor,
    )


def fetch_all_rows(cur, table: str) -> Tuple[List[str], List[tuple]]:
    cur.execute(f"SELECT * FROM `{table}`")
    rows = cur.fetchall()
    cols = [d[0] for d in cur.description] if cur.description else []
    return cols, list(rows)


def table_exists(cur, table: str) -> bool:
    cur.execute("SHOW TABLES LIKE %s", (table,))
    return cur.fetchone() is not None


def recreate_table_from_source(source, target, table: str) -> None:
    """用源库 DDL 在目标库 DROP 后重建，保证列与测试库一致（覆盖旧结构）。"""
    with source.cursor() as sc:
        sc.execute(f"SHOW CREATE TABLE `{table}`")
        row = sc.fetchone()
        if not row:
            raise RuntimeError(f"源库不存在表: {table}")
        ddl = row[1]
    with target.cursor() as tc:
        tc.execute("SET FOREIGN_KEY_CHECKS=0")
        tc.execute(f"DROP TABLE IF EXISTS `{table}`")
        tc.execute(ddl)
        tc.execute("SET FOREIGN_KEY_CHECKS=1")
    target.commit()


def dry_run(src: dict, dst: dict, tables: Sequence[str]) -> None:
    s = open_conn(**src)
    d = open_conn(**dst)
    try:
        with s.cursor() as sc, d.cursor() as dc:
            for t in tables:
                sc.execute(f"SELECT COUNT(*) FROM `{t}`")
                (n_src,) = sc.fetchone()
                if not table_exists(dc, t):
                    print(f"{t}: source_rows={n_src} target_rows=(表不存在)")
                else:
                    dc.execute(f"SELECT COUNT(*) FROM `{t}`")
                    (n_dst,) = dc.fetchone()
                    print(f"{t}: source_rows={n_src} target_rows={n_dst}")
    finally:
        s.close()
        d.close()


def sync_tables(src: dict, dst: dict, tables: Sequence[str]) -> None:
    source = open_conn(**src)
    target = open_conn(**dst)
    try:
        with source.cursor() as sc:
            for table in tables:
                recreate_table_from_source(source, target, table)
                print(f"++ {table}: 目标表已按源库 DDL 重建")
                cols, rows = fetch_all_rows(sc, table)
                if not cols:
                    raise RuntimeError(f"无法读取列: {table}")
                placeholders = ",".join(["%s"] * len(cols))
                col_sql = ",".join(f"`{c}`" for c in cols)
                insert_sql = f"INSERT INTO `{table}` ({col_sql}) VALUES ({placeholders})"
                with target.cursor() as tc:
                    if rows:
                        tc.executemany(insert_sql, rows)
                target.commit()
                print(f"OK {table}: copied {len(rows)} rows")
    finally:
        source.close()
        target.close()


def main() -> None:
    ap = argparse.ArgumentParser(description="Sync InfoGenie MySQL tables from test to production.")
    ap.add_argument(
        "--execute",
        action="store_true",
        help="执行同步（默认仅打印说明；不加此参数则退出码 1）",
    )
    ap.add_argument("--dry-run", action="store_true", help="仅统计源/目标行数，不写生产库")
    ap.add_argument(
        "--tables",
        type=str,
        default="",
        help="逗号分隔表名子集；默认同步全部 GORM 表",
    )
    args = ap.parse_args()

    src = conn_params_src()
    dst = conn_params_dst()

    tables: Tuple[str, ...]
    if args.tables.strip():
        tables = tuple(x.strip() for x in args.tables.split(",") if x.strip())
        bad = [t for t in tables if t not in TABLES]
        if bad:
            print(f"未知表名（不在白名单）: {bad}", file=sys.stderr)
            sys.exit(1)
    else:
        tables = TABLES

    print(f"SOURCE: {src['user']}@{src['host']}:{src['port']}/{src['database']}")
    print(f"TARGET: {dst['user']}@{dst['host']}:{dst['port']}/{dst['database']}")
    print(f"TABLES: {', '.join(tables)}")

    if args.dry_run:
        dry_run(src, dst, tables)
        return

    if not args.execute:
        print("未执行：请加 --dry-run 查看行数，或加 --execute 在确认备份后同步。", file=sys.stderr)
        sys.exit(1)

    sync_tables(src, dst, tables)
    print("完成。")


if __name__ == "__main__":
    main()