# -*- coding: utf-8 -*-
"""把 SQLite 库整体迁移进 MySQL（一次性工具，长期保留）。

用途：
  * 平台从 SQLite 迁到 MySQL 时搬数据
  * 从老备份（导出 zip 里的 users.db）恢复进 MySQL
  * 换环境（dev 库 ↔ 正式库）

用法示例：
    # 先干跑：只审计源库 + 建 schema，不搬数据
    python scripts/migrate_sqlite_to_mysql.py --sqlite data/users.db --env dev --dry-run

    # 正式迁（目标库已有数据会被整表替换）
    python scripts/migrate_sqlite_to_mysql.py --sqlite data/users.db --env dev --mode replace

    # 迁完只做校验（逐表 SHA-256 比对，不写任何数据）
    python scripts/migrate_sqlite_to_mysql.py --sqlite data/users.db --env dev --mode verify

    # 生产（必须显式二次确认）
    python scripts/migrate_sqlite_to_mysql.py --sqlite data/users.db --env prod --allow-prod --yes

目标库怎么定：默认读 .env / 环境变量的 DB_HOST/DB_PORT/DB_USER/DB_PASSWORD/DB_NAME
（也可用 --target-url 直接给完整连接串）。脚本**不会**碰源库（只读打开）。

安全设计：
  * 全程只读源库；目标库的写入在**单个事务**里，失败回滚，可放心重跑
  * 搬完做逐表行数 + 逐表全行 SHA-256 校验，不通过就报错
  * `--env prod` 必须额外 `--allow-prod`，交互终端还要手打 prod 确认
  * 目标库若已登记为别的环境（app_meta.deployment_env），除非 --force-env 否则拒绝
"""
import argparse
import hashlib
import json
import os
import sqlite3
import sys
import time
import uuid

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

# Windows 控制台默认 GBK，输出里带 ✔ 之类的符号会直接抛 UnicodeEncodeError
# （本脚本是运维在终端里手跑的，崩在最后一步会让人以为迁移失败）
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

from sqlalchemy import inspect as sa_inspect, text as sa_text  # noqa: E402


def _fail(msg):
    print("\n[错误] " + msg)
    sys.exit(2)


def _readonly_uri(path):
    from urllib.parse import quote
    return "file:{}?mode=ro".format(
        quote(os.path.abspath(path).replace("\\", "/"), safe="/:"))


def parse_args():
    ap = argparse.ArgumentParser(description="SQLite → MySQL 数据迁移")
    ap.add_argument("--sqlite", default=os.path.join(ROOT, "data", "users.db"),
                    help="源 SQLite 库路径（默认 data/users.db）")
    ap.add_argument("--env", choices=["dev", "prod"], required=True,
                    help="目标库属于哪个环境（写入 app_meta.deployment_env）")
    ap.add_argument("--mode", choices=["replace", "verify"], default="replace",
                    help="replace=建 schema 并搬数据；verify=只比对不写")
    ap.add_argument("--dry-run", action="store_true",
                    help="只审计源库并建 schema，不搬数据")
    ap.add_argument("--schema-only", action="store_true",
                    help="只建目标库的表结构，不搬数据")
    ap.add_argument("--target-url", default="",
                    help="目标连接串（默认取 .env / 环境变量里的 DB_*）")
    ap.add_argument("--allow-prod", action="store_true",
                    help="允许对生产库操作（--env prod 时必填）")
    ap.add_argument("--yes", action="store_true",
                    help="非交互环境下的最终确认（与 --allow-prod 一起用）")
    ap.add_argument("--force-env", action="store_true",
                    help="目标库已登记为别的环境时仍然继续（危险）")
    return ap.parse_args()


# ================== 1. 源库审计 ==================
def audit_source(src_path, tables):
    """校验源库可读、完整，并做列宽审计。返回 {表: {列: 最大长度}}。"""
    if not os.path.exists(src_path):
        _fail("源库不存在: " + src_path)
    con = sqlite3.connect(_readonly_uri(src_path), uri=True)
    con.text_factory = str
    try:
        integrity = con.execute("PRAGMA integrity_check").fetchone()[0]
        print(f"源库完整性: {integrity}")
        if integrity != "ok":
            _fail("源库完整性校验失败，请先修复源库再迁移")

        have = {r[0] for r in con.execute(
            "SELECT name FROM sqlite_master WHERE type='table'")}
        print(f"源库表（{len(have)}）: " + ", ".join(sorted(have)))
        missing = [t for t in tables if t not in have]
        if missing:
            print(f"  注意：模型里有但源库没有的表（迁移后为空）: {', '.join(missing)}")

        # 列宽审计：SQLite 不强制长度，MySQL 严格模式下超长会直接报错
        widths = {}
        for t in sorted(have):
            if t.startswith("sqlite_"):
                continue
            widths[t] = {}
            cols = [r[1] for r in con.execute('PRAGMA table_info("%s")' % t)]
            for c in cols:
                try:
                    n = con.execute(
                        'SELECT MAX(LENGTH("%s")) FROM "%s"' % (c, t)).fetchone()[0]
                except sqlite3.Error:
                    continue
                if n:
                    widths[t][c] = n
        return widths
    finally:
        con.close()


def check_widths(widths):
    """把源库实测列宽与模型声明比对，超限直接拦下（MySQL 严格模式会报错）。"""
    from core.models import db
    problems = []
    print("\n列宽审计（源库实测最大长度 / 模型上限）:")
    for table in db.metadata.sorted_tables:
        w = widths.get(table.name) or {}
        for col in table.columns:
            n = w.get(col.name)
            limit = getattr(col.type, "length", None)
            if n is None or not limit:
                continue
            flag = ""
            if n > limit:
                flag = "  <== 超出！"
                problems.append(f"{table.name}.{col.name}: {n} > {limit}")
            if n > limit * 0.6:
                print(f"   {table.name}.{col.name}: {n}/{limit}{flag}")
    if problems:
        _fail("以下列的实际数据超过模型声明长度，请先加长模型列宽或清理数据：\n  "
              + "\n  ".join(problems))
    print("   （只打印长度超过上限 60% 的列；全部在限内）")


# ================== 2. 目标库 schema ==================
def build_target_app(uri, env):
    """用目标连接串建一个最小 Flask app，复用平台自己的建表逻辑。"""
    os.environ["DEPLOY_ENV"] = env
    os.environ["DATABASE_URL"] = uri
    from flask import Flask
    from core.models import db  # noqa: E402

    app = Flask("migrate")
    app.config["SQLALCHEMY_DATABASE_URI"] = uri
    app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
    from core import db_config
    app.config["SQLALCHEMY_ENGINE_OPTIONS"] = db_config.engine_options(uri)
    db.init_app(app)
    with app.app_context():
        # 只做结构：建表 / 补列 / 唯一索引（跳过默认管理员与旧 JSON 迁移）
        from core.models import _ensure_unique_indexes, _sync_columns
        db.create_all()
        _sync_columns()
        _ensure_unique_indexes()
    return app


# ================== 3. 搬数据 + 校验 ==================
def _rows_from_sqlite(con, table, cols):
    names = [c.name for c in cols]
    sel = "SELECT {} FROM \"{}\"".format(
        ", ".join('"%s"' % n for n in names), table.name)
    cur = con.execute(sel)
    return names, cur.fetchall()


def _norm(v):
    """把驱动层返回的值统一成可跨库比较的形态。

    两个库的类型系统不一样，直接比 Python 对象会误报：
      * 布尔：SQLite 存 1/0（原始读出来是 int），走 ORM 读回来是 True/False
      * 整数值的浮点：SQLite 的 REAL 列里可能存着 7，MySQL 的 DOUBLE 读回来是 7.0
      * 文本：某些驱动可能给 bytes
    """
    if isinstance(v, bool):
        return int(v)
    if isinstance(v, float) and v.is_integer():
        return int(v)
    if isinstance(v, (bytes, bytearray)):
        return v.decode("utf-8", "replace")
    return v


def _canonical(names, rows):
    """把整表行规范化成一个字符串，用于跨库比对（行序无关：先排序再哈希）。"""
    packed = [json.dumps([_norm(r[i]) for i in range(len(names))],
                         ensure_ascii=False, default=str,
                         sort_keys=False) for r in rows]
    packed.sort()
    h = hashlib.sha256()
    for p in packed:
        h.update(p.encode("utf-8"))
        h.update(b"\n")
    return h.hexdigest()


def copy_and_verify(src_path, table_order, mode, chunk=500):
    """搬运并校验。返回 (逐表统计, 是否全部一致)。

    mode="replace"：DELETE 全表 + 插入 + 回读比对（单事务，失败整体回滚）
    mode="verify" ：只回读比对，不写任何数据
    """
    from core.models import db

    con = sqlite3.connect(_readonly_uri(src_path), uri=True)
    con.text_factory = str
    stats = []
    try:
        have = {r[0] for r in con.execute(
            "SELECT name FROM sqlite_master WHERE type='table'")}
        with db.engine.begin() as dst:          # 单事务：失败整体回滚
            if mode == "replace":
                for table in reversed(table_order):
                    dst.execute(table.delete())
            for table in table_order:
                if table.name not in have:
                    stats.append((table.name, 0, "源库无此表", "-"))
                    continue
                db_cols = {c["name"] for c in sa_inspect(dst).get_columns(table.name)}
                cols = [c for c in table.columns if c.name in db_cols]
                names, rows = _rows_from_sqlite(con, table, cols)
                src_hash = _canonical(names, rows)
                if mode == "replace":
                    for i in range(0, len(rows), chunk):
                        batch = [dict(zip(names, r)) for r in rows[i:i + chunk]]
                        dst.execute(table.insert(), batch)
                stats.append((table.name, len(rows), src_hash))
    finally:
        con.close()

    # 回读目标库重算（verify 模式只有这一步）
    with db.engine.connect() as dst:
        out = []
        for (name, srows, shash) in stats:
            table = next((t for t in table_order if t.name == name), None)
            if table is None or shash == "-":
                out.append((name, srows, "-", "源库无此表", "-"))
                continue
            db_cols = {c["name"] for c in sa_inspect(dst).get_columns(name)}
            names = [c.name for c in table.columns if c.name in db_cols]
            # 用原生 SQL 回读：走 ORM/Core 的 typed select 会做类型转换
            # （Boolean→True/False），与源库驱动层读出来的 1/0 对不上
            prep = dst.dialect.identifier_preparer
            sel = "SELECT {} FROM {}".format(
                ", ".join(prep.quote(n) for n in names), prep.quote(name))
            rows = [tuple(r) for r in dst.execute(sa_text(sel)).fetchall()]
            thash = _canonical(names, rows)
            same = (len(rows) == srows) and (thash == shash)
            out.append((name, srows, len(rows), "一致" if same else "不一致", shash))
    return out, all(r[3] == "一致" or r[3] == "源库无此表" for r in out)


def print_stats(stats):
    print("\n%-20s %6s %6s  %-10s %s" % ("表", "源", "目标", "结果", "SHA-256(前12)"))
    for (name, srows, trows, verdict, shash) in stats:
        print("%-20s %6s %6s  %-10s %s" % (name, srows, trows, verdict, (shash or "-")[:12]))


# ================== 主流程 ==================
def main():
    args = parse_args()

    if args.env == "prod":
        if not args.allow_prod:
            _fail("--env prod 必须同时给 --allow-prod（这条命令会写生产库）")
        if sys.stdin.isatty() and not args.yes:
            print("即将对【生产库】执行迁移。")
            if input("请输入 prod 确认: ").strip() != "prod":
                _fail("未确认，已中止")
        elif not args.yes:
            _fail("非交互环境下请追加 --yes 明确确认")

    from core import db_config
    uri = args.target_url or db_config.build_db_uri()
    if db_config.is_sqlite(uri):
        _fail("目标仍是 SQLite —— 请在 .env 里配好 DB_HOST/DB_USER/DB_PASSWORD/DB_NAME，"
              "或用 --target-url 指定 MySQL 连接串")
    print(f"源库   : {args.sqlite}")
    print(f"目标库 : {db_config.describe_target(uri)}   (环境 {args.env})")

    from core.models import db
    tables = list(db.metadata.sorted_tables)
    print(f"模型表（{len(tables)}）: " + ", ".join(t.name for t in tables))

    widths = audit_source(args.sqlite, [t.name for t in tables])
    check_widths(widths)

    if args.mode == "verify":
        # 不建 schema、不写数据，直接比对（用于迁移后复验）
        os.environ["DEPLOY_ENV"] = args.env
        os.environ["DATABASE_URL"] = uri
        from flask import Flask
        app = Flask("verify")
        app.config["SQLALCHEMY_DATABASE_URI"] = uri
        app.config["SQLALCHEMY_ENGINE_OPTIONS"] = db_config.engine_options(uri)
        db.init_app(app)
        with app.app_context():
            stats, ok = copy_and_verify(args.sqlite, tables, "verify")
        print_stats(stats)
        print("\n校验结果: " + ("全部一致 OK" if ok else "存在不一致 FAIL"))
        sys.exit(0 if ok else 1)

    # 目标库环境标签检查（防止误把 dev 数据灌进生产库）
    app = build_target_app(uri, args.env)
    with app.app_context():
        recorded = db_config.meta_get("deployment_env")
        if recorded and recorded != args.env and not args.force_env:
            _fail(f"目标库已登记为 {recorded} 环境，与 --env {args.env} 不符。"
                  f"确认无误请加 --force-env")
        if not recorded:
            print(f"目标库首次使用，将登记为 {args.env} 环境")

        if args.schema_only or args.dry_run:
            print("\n[dry-run] schema 已建好，未搬数据。")
            if args.dry_run:
                print("  去掉 --dry-run 即正式搬运。")
            return

        print("\n开始搬运（单事务，失败自动回滚）…")
        t0 = time.time()
        stats, ok = copy_and_verify(args.sqlite, tables, "replace")
        print_stats(stats)
        if not ok:
            _fail("搬运后校验不一致 —— 事务已回滚，目标库保持原样。请把上面的表格发给我排查")
        print(f"\n搬运完成，用时 {time.time() - t0:.1f}s，逐表校验一致 OK")

        # 4) 写库标签与版本
        src = sqlite3.connect(_readonly_uri(args.sqlite), uri=True)
        try:
            src_id = None
            try:
                row = src.execute(
                    "SELECT value FROM app_meta WHERE key='deployment_id'").fetchone()
                src_id = row[0] if row else None
            except sqlite3.Error:
                pass
        finally:
            src.close()
        db_config.meta_set("deployment_env", args.env)
        db_config.meta_set("deployment_id", src_id or uuid.uuid4().hex)
        db_config.meta_set("deployment_claimed_at",
                           time.strftime("%Y-%m-%d %H:%M:%S"))
        from core.models import CURRENT_SCHEMA_VERSION
        db_config.meta_set("schema_version", str(CURRENT_SCHEMA_VERSION))
        print(f"已登记库环境标签: {args.env}（部署标识 {db_config.meta_get('deployment_id')}）")
        print("\n下一步：把 .env 的 DEPLOY_ENV/DB_* 配好，重启 web_server 即连到新库。")


if __name__ == "__main__":
    main()