# +======================================================================
# | DjangoAdmin 企业数字化开发平台 [ 赋能开发者，助力企业发展 ]
# +----------------------------------------------------------------------
# | Copyright (c) 2021-2026 小蚂蚁云团队 版权所有
# +----------------------------------------------------------------------
# | 开源协议: Apache-2.0
# | 重要声明: 未经书面许可，不得移除或修改本框架相关版权标识。
# +----------------------------------------------------------------------
# | 官方网站: https://www.djangoadmin.cn
# +----------------------------------------------------------------------
# | 研发团队: 小蚂蚁云团队 荣誉出品
# +======================================================================
# | 版权和免责声明:
# | 一、知识产权声明
# | 本团队对该软件框架产品（以下简称“本产品”）拥有完整知识产权，包括但不限于商标权、专利权、
# | 著作权、商业秘密等。前述权利均受中华人民共和国及相关国际法律法规保护。未经本团队书面授权，
# | 任何个人或单位不得基于本产品本身申请任何形式的知识产权。
# | 二、使用限制与禁止条款
# | 您理解并同意，禁止将本产品用于以下任何行为：
# | 1. 任何违反中华人民共和国法律法规的用途；
# | 2. 任何侵害他人合法权益（包括但不限于人身权、财产权、知识产权等）的行为；
# | 3. 任何危害国家安全、公共利益及社会道德风尚的行为。
# | 三、责任豁免
# | 因使用本产品而产生的任何损失或法律纠纷（包括但不限于直接损失、间接损失、附带损失、衍生损失
# | 以及因违约、侵权等引发的责任），本团队均不承担任何法律责任。若因违规使用本产品导致任何第三
# | 方提起索赔或行政机关进行调查，您应独立承担全部责任。
# | 四、违约与法律责任
# | 如发现任何违反本声明的行为，本团队有权：
# | 1. 无条件配合公安机关及其他有权机关的调查取证工作；
# | 2. 依法追究相关方的民事责任，并要求赔偿全部损失；
# | 3. 保留进一步采取刑事控告等法律措施的权利。
# | 五、许可范围说明
# | 1、被授权方应视本产品源代码为核心机密信息，负有最高标准的保密与保管义务。禁止以任何形式（包
# | 括但不限于直接或间接的分享、复制、传播、分发、转售、出租、开源、反向工程、反编译、反汇编或通
# | 过网络服务提供给任何第三方）泄露、披露或处置源代码的全部或部分内容。任何未经授权的复制、使用、
# | 传播行为均构成对本团队知识产权的严重侵权，本团队有权依法追究其全部法律责任，并要求赔偿因此造
# | 成的一切经济损失及商誉损害。
# | 2、本产品仅限于公司或个人在法律允许范围内，用于内部且合规的软件产品研发。具体权利义务以《软件
# | 许可协议》附件为准。
# +======================================================================

# +======================================================================
# | 模块: 跨库数据迁移脚本（Django版本）
# | 说明: 把源库的数据整体搬到目标库（MySQL → PostgreSQL / SQL Server / SQLite / Oracle）
# +======================================================================

"""
跨库数据迁移脚本（Django版本）
=============================================================

把「源库」的数据整体搬到「目标库」：
- 目标库结构用 SQLAlchemy ORM 模型 metadata 生成（Base.metadata.create_all，
  方言可移植，模型即 head）；
- 数据按拓扑序逐表 SELECT ... ORDER BY id，显式带 id 批量插入，保住
  无外键裸列引用关系（user.dept_id、role_menu.role_id/menu_id 等）；
- 最后按方言重置自增序列（PG setval / MSSQL IDENTITY_INSERT + DBCC
  CHECKIDENT / Oracle 为每表重建 sequence + BEFORE INSERT 触发器；
  MySQL / SQLite 显式插 id 后自动跟随，无需处理）。

与 init_db.py 的分工：
- init_db.py    单库初始化（Django migrate）
- migrate_db.py 双库迁移（源库 → 目标库，CLI 显式参数）

用法：
    python scripts/migrate_db.py \
        --src-driver=mysql --src-host=127.0.0.1 --src-port=3306 \
        --src-db=djangoadmin.django.elevue --src-user=root --src-pass=root \
        --dst-driver=postgresql --dst-host=127.0.0.1 --dst-port=5432 \
        --dst-db=djangoadmin.django.elevue --dst-user=root --dst-pass=root \
        --drop-target-first

先干跑核对计划：加 --dry-run（只连源、逐表数行数、打印计划，不碰目标库）。
"""

import argparse
import os
import sys
import time
from pathlib import Path
from urllib.parse import quote_plus

# 设置项目根目录
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

# Windows 控制台编码处理
if sys.stdout and hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')

from sqlalchemy import create_engine, inspect, select, text, Column, Integer, String, DateTime, Boolean, Text, SmallInteger, BigInteger
from sqlalchemy.pool import NullPool
from sqlalchemy.orm import declarative_base

# 加载 Django 配置
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'application.settings')

# 创建 SQLAlchemy Base（方言可移植的模型定义）
Base = declarative_base()

# 各驱动默认端口
_DEFAULT_PORTS = {"mysql": 3306, "postgresql": 5432, "mssql": 1433, "oracle": 1521, "sqlite": None}

# 迁移顺序：基础/被引用表在前，依赖表在后
# Django内置表（无前缀）
_DJANGO_BUILTIN_BASE = [
    "auth_permission",
    "auth_group",
    "auth_user",
    "django_content_type",
    "django_migrations",
]
_DJANGO_BUILTIN_DEP = [
    "auth_group_permissions",
    "auth_user_groups",
    "auth_user_user_permissions",
    "captcha_captchastore",
    "django_session",
    "django_admin_log",
]

# 项目业务表（带前缀）
_BASE_TABLES = [
    "dept", "menu", "role", "level", "position", "city", "category",
    "dict", "config", "file_template", "param", "job", "link",
]
_DEP_TABLES = [
    "role_menu", "dict_item", "config_item", "job_log",
    "user", "user_role", "article", "notice", "operation_log", "login_log",
]

# 所有需要迁移的表（35张）
_ALL_TABLES = (
    _DJANGO_BUILTIN_BASE +
    _DJANGO_BUILTIN_DEP +
    [t for t in _BASE_TABLES] +
    [t for t in _DEP_TABLES]
)


# ============================================================
# Django 模型映射为 SQLAlchemy ORM 模型（方言可移植）
# ============================================================
def load_django_models(prefix: str):
    """从 Django 配置加载模型信息，动态创建 SQLAlchemy ORM 模型

    关键设计：使用 SQLAlchemy 原生类型（Integer, String, DateTime 等），
    这些类型在所有数据库方言中都有对应实现，保证 create_all 方言可移植。

    迁移范围：所有 Django 模型（包括内置表 auth_*、django_*、captcha_* 等）
    特别处理：Django多对多关系表和迁移记录表需要手动添加
    """
    import django
    django.setup()

    from django.apps import apps

    # 获取所有模型（包括Django内置表和业务表）
    models = []
    for model in apps.get_models():
        table_name = model._meta.db_table
        models.append((table_name, model))

    # 手动添加Django多对多关系表和迁移记录表（这些表在模型中没有直接定义）
    # auth_group_permissions: Group和Permission的多对多关系表
    # auth_user_groups: User和Group的多对多关系表
    # auth_user_user_permissions: User和Permission的多对多关系表
    # django_migrations: Django迁移记录表

    # 按拓扑序排序
    def sort_key(item):
        table_name = item[0]

        # Django内置表（无前缀）
        if not table_name.startswith(prefix):
            if table_name in _DJANGO_BUILTIN_BASE:
                return (0, _DJANGO_BUILTIN_BASE.index(table_name))
            elif table_name in _DJANGO_BUILTIN_DEP:
                return (1, _DJANGO_BUILTIN_DEP.index(table_name))
            else:
                return (2, 0)
        # 项目业务表（带前缀）
        else:
            name = table_name[len(prefix):]  # 去掉前缀
            if name in _BASE_TABLES:
                return (3, _BASE_TABLES.index(name))
            elif name in _DEP_TABLES:
                return (4, _DEP_TABLES.index(name))
            else:
                return (5, 0)

    models.sort(key=sort_key)

    # 动态创建 SQLAlchemy 模型
    for table_name, django_model in models:
        # 获取字段定义
        columns = []
        for field in django_model._meta.get_fields():
            if hasattr(field, 'column'):
                col_type = _get_sa_type(field)
                if col_type:
                    col_kwargs = {
                        'primary_key': field.primary_key if hasattr(field, 'primary_key') else False,
                        'nullable': field.null if hasattr(field, 'null') else True,
                    }
                    # 使用数据库列名（对于ForeignKey字段，column是带_id后缀的）
                    col_name = field.column if hasattr(field, 'column') else field.name
                    # 确保col_name不为None
                    if col_name:
                        columns.append(Column(col_name, col_type, **col_kwargs))

        # 创建表类
        if columns:
            type(table_name, (Base,), {
                '__tablename__': table_name,
                '__table_args__': {'extend_existing': True},
                **{col.name: col for col in columns}
            })

    # 手动添加Django多对多关系表和迁移记录表（这些表在模型中没有直接定义）
    # auth_group_permissions: Group和Permission的多对多关系表
    if 'auth_group_permissions' not in Base.metadata.tables:
        type('auth_group_permissions', (Base,), {
            '__tablename__': 'auth_group_permissions',
            '__table_args__': {'extend_existing': True},
            'id': Column(Integer, primary_key=True),
            'group_id': Column(Integer, nullable=False),
            'permission_id': Column(Integer, nullable=False),
        })

    # auth_user_groups: User和Group的多对多关系表
    if 'auth_user_groups' not in Base.metadata.tables:
        type('auth_user_groups', (Base,), {
            '__tablename__': 'auth_user_groups',
            '__table_args__': {'extend_existing': True},
            'id': Column(Integer, primary_key=True),
            'user_id': Column(Integer, nullable=False),
            'group_id': Column(Integer, nullable=False),
        })

    # auth_user_user_permissions: User和Permission的多对多关系表
    if 'auth_user_user_permissions' not in Base.metadata.tables:
        type('auth_user_user_permissions', (Base,), {
            '__tablename__': 'auth_user_user_permissions',
            '__table_args__': {'extend_existing': True},
            'id': Column(Integer, primary_key=True),
            'user_id': Column(Integer, nullable=False),
            'permission_id': Column(Integer, nullable=False),
        })

    # django_migrations: Django迁移记录表
    if 'django_migrations' not in Base.metadata.tables:
        type('django_migrations', (Base,), {
            '__tablename__': 'django_migrations',
            '__table_args__': {'extend_existing': True},
            'id': Column(Integer, primary_key=True),
            'app': Column(String(255), nullable=False),
            'name': Column(String(255), nullable=False),
            'applied': Column(DateTime, nullable=False),
        })


def _get_sa_type(field):
    """将 Django 字段类型转换为 SQLAlchemy 类型"""
    from django.db import models as django_models

    field_type = type(field)

    # 带长度的类型需要特殊处理
    if issubclass(field_type, django_models.CharField):
        max_length = getattr(field, 'max_length', 255) or 255
        return String(max_length)

    # 映射表（不带长度的类型）
    type_mapping = {
        django_models.AutoField: Integer,
        django_models.IntegerField: Integer,
        django_models.SmallIntegerField: SmallInteger,
        django_models.BigIntegerField: BigInteger,
        django_models.TextField: Text,
        django_models.BooleanField: Boolean,
        django_models.DateTimeField: DateTime,
        django_models.DateField: DateTime,
        django_models.TimeField: DateTime,
        django_models.DecimalField: String(50),  # 用String存储，避免精度问题
        django_models.FloatField: String(50),
        django_models.EmailField: String(254),
        django_models.URLField: String(200),
        django_models.SlugField: String(50),
        django_models.IPAddressField: String(15),
        django_models.GenericIPAddressField: String(39),
        django_models.FileField: String(100),
        django_models.ImageField: String(100),
    }

    # 查找匹配的类型
    for django_type, sa_type in type_mapping.items():
        if issubclass(field_type, django_type):
            return sa_type if not callable(sa_type) else sa_type

    # 默认返回String
    return String(255)


# ============================================================
# CLI 参数解析
# ============================================================
def build_parser() -> argparse.ArgumentParser:
    """CLI 参数：源/目标两端连接显式给出"""
    p = argparse.ArgumentParser(description="跨库数据迁移：源库 → 目标库（Django版本）")

    def add_side(side: str, default_user: str):
        p.add_argument(f"--{side}-driver", choices=("mysql", "postgresql", "mssql", "oracle", "sqlite"), default="mysql",
                       help="驱动（oracle 时 --{side}-db 即 service_name；sqlite 时即数据库文件路径）")
        p.add_argument(f"--{side}-host", default="127.0.0.1")
        p.add_argument(f"--{side}-port", type=int, default=None,
                       help="端口，缺省按驱动默认（mysql 3306 / pg 5432 / mssql 1433 / oracle 1521）")
        p.add_argument(f"--{side}-db", required=True, help="数据库名（sqlite 时为文件路径）")
        p.add_argument(f"--{side}-user", default=default_user)
        p.add_argument(f"--{side}-pass", default="",
                       help=f"口令，留空回落环境变量 {side.upper()}_DB_PASS（避免进 shell 历史）")

    add_side("src", "root")
    add_side("dst", "postgres")

    p.add_argument("--drop-target-first", action="store_true", help="先删目标库再重建（默认关）")
    p.add_argument("--truncate-target", action="store_true", help="先清空目标表数据（默认关，与 --drop-target-first 互斥）")
    p.add_argument("--dry-run", action="store_true", help="只连源、逐表数行数、打印计划，不碰目标库")
    p.add_argument("--chunk-size", type=int, default=1000, help="每批插入行数（默认 1000）")
    p.add_argument("--only-tables", default="", help="逗号分隔的表名子集（全名或去掉前缀的短名），续传用")
    p.add_argument("--skip-existing", action="store_true", help="目标已等于源行数的表跳过")
    p.add_argument("--skip-verify", action="store_true", help="跳过迁移后的逐表行数校验")
    p.add_argument("--single-transaction", action="store_true", help="全部表包进一个大事务（默认每表一个事务）")
    p.add_argument("--prefix", default=None, help="数据表前缀（默认从 Django 配置读取）")

    return p


# ============================================================
# 配置加载
# ============================================================
def load_django_config():
    """从 Django 配置中读取数据库连接信息"""
    try:
        from config import env
        return {
            'engine': getattr(env, 'DATABASE_ENGINE', 'django.db.backends.mysql'),
            'name': getattr(env, 'DATABASE_NAME', 'djangoadmin.django.elevue'),
            'host': getattr(env, 'DATABASE_HOST', '127.0.0.1'),
            'port': getattr(env, 'DATABASE_PORT', 3306),
            'user': getattr(env, 'DATABASE_USER', 'root'),
            'password': getattr(env, 'DATABASE_PASSWORD', 'root'),
            'prefix': getattr(env, 'DATABASE_PREFIX', 'django_'),
        }
    except ImportError:
        return None


def django_engine_to_driver(engine_str: str) -> str:
    """将 Django 数据库引擎转换为 SQLAlchemy 驱动名"""
    mapping = {
        'django.db.backends.mysql': 'mysql',
        'django.db.backends.postgresql': 'postgresql',
        'django.db.backends.sqlite3': 'sqlite',
        'django.db.backends.oracle': 'oracle',
        'mssql': 'mssql',
    }
    return mapping.get(engine_str, 'mysql')


# ============================================================
# URL / 引擎
# ============================================================
def _normalize_driver(driver: str) -> str:
    """标准化驱动名称"""
    aliases = {
        'postgres': 'postgresql',
        'psycopg2': 'postgresql',
        'psycopg': 'postgresql',
        'mysqlclient': 'mysql',
        'pymysql': 'mysql',
        'sqlite3': 'sqlite',
        'mssql': 'mssql',
        'sqlserver': 'mssql',
        'oracle': 'oracle',
    }
    return aliases.get(driver.lower(), driver.lower())


def _sqlite_url(database: str) -> str:
    """构建 SQLite 连接 URL"""
    path = Path(database).expanduser().resolve()
    return f"sqlite:///{path}"


def make_url(driver: str, host: str, port: int, database: str, username: str, password: str) -> str:
    """构建数据库连接 URL"""
    drv = _normalize_driver(driver)
    if drv == 'sqlite':
        return _sqlite_url(database)

    base = quote_plus(username) + ':' + quote_plus(password) + '@' + host + ':' + str(port)

    if drv == 'oracle':
        return 'oracle+oracledb://' + base + '/?service_name=' + quote_plus(database)

    base = base + '/' + database

    if drv == 'mysql':
        return 'mysql+pymysql://' + base + '?charset=utf8mb4'
    if drv == 'postgresql':
        return 'postgresql+psycopg://' + base

    # mssql
    return 'mssql+pymssql://' + base + '?charset=utf8'


def make_engine(driver: str, host: str, port: int, database: str, username: str, password: str):
    """迁移专用引擎：NullPool，避免连接池把目标/源的会话状态串起来"""
    return create_engine(make_url(driver, host, port, database, username, password), poolclass=NullPool)


def _resolve_port(driver: str, port: int | None) -> int | None:
    """解析端口号，如果未指定则使用默认端口"""
    return port if port is not None else _DEFAULT_PORTS.get(driver)


def _fmt_duration(sec: float) -> str:
    """耗时格式化：<1s 显示毫秒，<60s 显示秒，否则分+秒"""
    if sec < 1:
        return f"{sec * 1000:.0f}ms"
    if sec < 60:
        return f"{sec:.1f}s"
    return f"{int(sec // 60)}m{int(sec % 60)}s"


_BUSY = "[·]"
_DONE = "[✓]"


def _phase_start(n: int, total: int, desc: str, hint: bool = True):
    """阶段开始横幅"""
    tail = "（可能耗时较长，请勿中断）" if hint else ""
    print(f"{_BUSY} [{n}/{total}] {desc}...{tail}")


def _phase_done(n: int, total: int, desc: str, t0: float, extra: str = ""):
    """阶段完成横幅"""
    print(f"{_DONE} [{n}/{total}] {desc} 完成（耗时 {_fmt_duration(time.perf_counter() - t0)}）{extra}")


# ============================================================
# 建目标库
# ============================================================
def create_database_if_missing(driver: str, host: str, port: int, database: str,
                               username: str, password: str, drop_first: bool = False,
                               tables: tuple = ()):
    """按方言连维护库建目标库（不存在才建）。sqlite 无建库概念，只确保父目录并可选删旧文件"""
    if driver == 'sqlite':
        file = Path(database).expanduser()
        print(f"  · 检查/准备 sqlite 文件 {file}...")
        if file.parent.as_posix() != '.':
            file.parent.mkdir(parents=True, exist_ok=True)
        if drop_first and file.exists():
            print(f"  · 删除旧 sqlite 文件 {file}...")
            file.unlink()
        return

    if driver == 'mysql':
        import pymysql
        print("  · 正在连接 MySQL（无库连接）...")
        conn = pymysql.connect(host=host, port=int(port), user=username, password=password, autocommit=True)
        cur = conn.cursor()
        if drop_first:
            print(f"  · 正在删除旧目标库 `{database}`...")
            cur.execute(f"DROP DATABASE IF EXISTS `{database}`")
        print(f"  · 正在创建目标库 `{database}`（utf8mb4）...")
        cur.execute(
            f"CREATE DATABASE IF NOT EXISTS `{database}` "
            f"DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci"
        )
        cur.close()
        conn.close()
        return

    if driver == 'postgresql':
        import psycopg
        print("  · 正在连接 postgres 维护库...")
        conn = psycopg.connect(
            host=host, port=int(port), user=username, password=password,
            dbname="postgres", autocommit=True,
        )
        with conn.cursor() as cur:
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (database,))
            exists = cur.fetchone() is not None
            if drop_first and exists:
                print(f"  · 正在删除旧目标库 \"{database}\"（WITH FORCE）...")
                cur.execute(f'DROP DATABASE "{database}" WITH (FORCE)')
                exists = False
            if not exists:
                print(f"  · 正在创建目标库 \"{database}\"（UTF8）...")
                cur.execute(f'CREATE DATABASE "{database}" ENCODING \'UTF8\' TEMPLATE template0')
        conn.close()
        return

    if driver == 'oracle':
        import oracledb
        print(f"  · 正在连接 Oracle 服务 {database}（schema={username}）...")
        conn = oracledb.connect(user=username, password=password, dsn=f"{host}:{port}/{database}")
        cur = conn.cursor()
        cur.execute("SELECT 1 FROM dual")
        cur.fetchone()
        if drop_first:
            print(f"  · drop-first：清理 {len(tables)} 张旧项目表（DROP TABLE ... PURGE）...")
            for t in tables:
                try:
                    cur.execute(f"DROP TABLE {t} PURGE")
                except Exception:
                    pass
        conn.close()
        return

    # mssql
    import pymssql
    print("  · 正在连接 SQL Server（master 维护库）...")
    conn = pymssql.connect(
        server=host, user=username, password=password,
        port=int(port), database="master", autocommit=True,
    )
    cur = conn.cursor()
    if drop_first:
        print(f"  · 正在删除旧目标库 [{database}]（SINGLE_USER + DROP）...")
        cur.execute(
            f"IF DB_ID(N'{database}') IS NOT NULL "
            f"ALTER DATABASE [{database}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE"
        )
        cur.execute(f"IF DB_ID(N'{database}') IS NOT NULL DROP DATABASE [{database}]")
    print(f"  · 正在创建目标库 [{database}]...")
    cur.execute(f"IF DB_ID(N'{database}') IS NULL CREATE DATABASE [{database}]")
    cur.close()
    conn.close()


# ============================================================
# 迁移顺序
# ============================================================
def ordered_tables(prefix: str):
    """按拓扑序返回要迁移的 Table 对象

    排序规则：
    1. Django内置基础表（auth_permission, auth_group, auth_user 等）
    2. Django内置依赖表（auth_group_permissions, django_session 等）
    3. 项目业务基础表（dept, menu, role 等）
    4. 项目业务依赖表（role_menu, user_role 等）
    5. 其他表
    """
    out, seen = [], set()

    # 1. Django内置基础表
    for name in _DJANGO_BUILTIN_BASE:
        tbl = Base.metadata.tables.get(name)
        if tbl is not None:
            out.append(tbl)
            seen.add(tbl.name)

    # 2. Django内置依赖表
    for name in _DJANGO_BUILTIN_DEP:
        tbl = Base.metadata.tables.get(name)
        if tbl is not None:
            out.append(tbl)
            seen.add(tbl.name)

    # 3. 项目业务基础表
    for name in _BASE_TABLES:
        tbl = Base.metadata.tables.get(prefix + name)
        if tbl is not None:
            out.append(tbl)
            seen.add(tbl.name)

    # 4. 项目业务依赖表
    for name in _DEP_TABLES:
        tbl = Base.metadata.tables.get(prefix + name)
        if tbl is not None:
            out.append(tbl)
            seen.add(tbl.name)

    # 5. 其他表
    for tbl in Base.metadata.tables.values():
        if tbl.name not in seen:
            out.append(tbl)

    return out


# ============================================================
# 泵数据
# ============================================================
def sanitize(payload: list[dict], table) -> list[dict]:
    """把驱动读回的零日期/非法日期字符串统一置 None（防御性兜底）"""
    for row in payload:
        for col in table.columns:
            v = row.get(col.name)
            if isinstance(v, str) and v.startswith('0000-00-00'):
                row[col.name] = None
    return payload


def pump_table(src_engine, dst_engine, table, chunk_size: int, driver: str) -> int:
    """逐 chunk 批量插入一张表，显式带 id 以保住裸列引用关系

    对于SQL Server，只有当表有IDENTITY列时才执行SET IDENTITY_INSERT
    对于大表，使用更小的批次并频繁提交，避免事务日志溢出
    """
    pk = next(iter(table.primary_key.columns))
    sel = select(*table.columns).order_by(pk)

    # 检查是否有IDENTITY列（自增列）
    # 只有当主键列是AutoIncrement时才需要SET IDENTITY_INSERT
    has_identity = False
    if driver == 'mssql':
        # 检查主键列是否是自增的
        pk_col = pk
        if hasattr(pk_col, 'autoincrement') and pk_col.autoincrement:
            has_identity = True
        # 也检查列名是否为id（Django的惯例）
        elif pk_col.name == 'id' and hasattr(pk_col, 'type') and isinstance(pk_col.type, Integer):
            # 对于id列，假设是自增的（Django惯例）
            has_identity = True

    # 对于SQL Server，使用更小的批次大小
    if driver == 'mssql':
        # 对于大表，使用更小的批次
        chunk_size = min(chunk_size, 500)

    total = 0
    t0 = time.perf_counter()

    # 对于SQL Server，使用独立的事务来处理每个批次
    # 避免事务日志溢出
    if driver == 'mssql':
        # SQL Server：每个批次使用独立的事务
        with src_engine.connect() as src:
            result = src.execute(sel)
            while True:
                rows = result.fetchmany(chunk_size)
                if not rows:
                    break
                payload = sanitize([dict(r._mapping) for r in rows], table)

                # 每个批次使用独立的事务
                with dst_engine.begin() as dst:
                    if has_identity:
                        try:
                            dst.execute(text(f"SET IDENTITY_INSERT [{table.name}] ON"))
                        except Exception:
                            has_identity = False

                    dst.execute(table.insert(), payload)

                    if has_identity:
                        try:
                            dst.execute(text(f"SET IDENTITY_INSERT [{table.name}] OFF"))
                        except Exception:
                            pass

                total += len(payload)
                print(f"  [{table.name}] +{len(payload)}（累计 {total}）")
    else:
        # 其他数据库：使用单个事务
        with dst_engine.begin() as dst:
            with src_engine.connect() as src:
                result = src.execute(sel)
                while True:
                    rows = result.fetchmany(chunk_size)
                    if not rows:
                        break
                    payload = sanitize([dict(r._mapping) for r in rows], table)
                    dst.execute(table.insert(), payload)
                    total += len(payload)
                    print(f"  [{table.name}] +{len(payload)}（累计 {total}）")

    print(f"  [✓ {table.name}] 共 {total} 行（耗时 {_fmt_duration(time.perf_counter() - t0)}）")
    return total


# ============================================================
# 自增序列重置
# ============================================================
def pg_reset_sql(table_name: str) -> str:
    """PostgreSQL：把 SERIAL/IDENTITY 序列推进到 MAX(id)"""
    return (
        f"SELECT setval(pg_get_serial_sequence('{table_name}', 'id'), "
        f"COALESCE((SELECT MAX(id) FROM \"{table_name}\"), 1), "
        f"(SELECT MAX(id) FROM \"{table_name}\") IS NOT NULL)"
    )


def mssql_reset_sql(table_name: str, max_id) -> str:
    """SQL Server：RESEED 必须是字面量，先查 MAX(id) 再拼"""
    return f"DBCC CHECKIDENT ('{table_name}', RESEED, {max_id or 0})"


def oracle_sequence_sql(table_name: str, start_with: int) -> str:
    """Oracle：为 INTEGER 主键建 sequence"""
    return f"CREATE SEQUENCE {table_name}_id_seq START WITH {start_with} INCREMENT BY 1 NOCACHE"


def oracle_trigger_sql(table_name: str) -> str:
    """Oracle：BEFORE INSERT 触发器，id 为 NULL 时从 sequence 取 NEXTVAL"""
    seq = f"{table_name}_id_seq"
    trg = f"{table_name}_id_trg"
    return (
        f"CREATE OR REPLACE TRIGGER {trg}\n"
        f"BEFORE INSERT ON {table_name} FOR EACH ROW\n"
        f"WHEN (NEW.id IS NULL)\n"
        f"BEGIN SELECT {seq}.NEXTVAL INTO :NEW.id FROM dual; END;"
    )


def ensure_oracle_autoincrement(conn, table_name: str):
    """Oracle 迁移后补自增机制"""
    max_id = conn.execute(text(f"SELECT COALESCE(MAX(id), 0) FROM {table_name}")).scalar()
    seq = f"{table_name}_id_seq"
    exists = conn.execute(text(
        "SELECT COUNT(*) FROM user_sequences WHERE sequence_name = UPPER(:n)"
    ), {"n": seq}).scalar()
    if exists:
        conn.execute(text(f"DROP SEQUENCE {seq}"))
    conn.execute(text(oracle_sequence_sql(table_name, max_id + 1)))
    conn.execute(text(oracle_trigger_sql(table_name)))


def reset_sequence(driver: str, conn, table):
    """迁移后重置自增序列

    跳过没有id字段的表（如django_session使用session_key作为主键）
    """
    # 检查表是否有id字段
    has_id_column = any(col.name == 'id' for col in table.columns)
    if not has_id_column:
        return  # 跳过没有id字段的表

    if driver == 'postgresql':
        conn.execute(text(pg_reset_sql(table.name)))
    elif driver == 'mssql':
        max_id = conn.execute(text(f"SELECT MAX(id) FROM [{table.name}]")).scalar()
        conn.execute(text(mssql_reset_sql(table.name, max_id)))
    elif driver == 'oracle':
        ensure_oracle_autoincrement(conn, table.name)


# ============================================================
# 校验
# ============================================================
def verify_counts(src, dst, tables) -> list[tuple[str, int, int]]:
    """逐表 count 对比，返回不一致的表列表"""
    failed = []
    for t in tables:
        s = src.execute(text(f"SELECT COUNT(*) FROM {t.name}")).scalar()
        d = dst.execute(text(f"SELECT COUNT(*) FROM {t.name}")).scalar()
        mark = "PASS" if s == d else "FAIL"
        print(f"  [{t.name}] source={s} target={d} {mark}")
        if s != d:
            failed.append((t.name, s, d))
    return failed


def _existing_equal(src, dst_engine, table) -> bool:
    """检查源表和目标表的行数是否相等

    对于SQL Server，使用独立的连接来执行查询，避免事务状态问题
    """
    s = src.execute(text(f"SELECT COUNT(*) FROM {table.name}")).scalar()
    # 使用独立的连接来查询目标表
    with dst_engine.connect() as dst:
        d = dst.execute(text(f"SELECT COUNT(*) FROM {table.name}")).scalar()
    return s == d


# ============================================================
# 主流程
# ============================================================
def _dry_run(src_engine, tables, dst_driver: str, dst_db: str, t_start: float) -> int:
    print("=" * 62)
    print(f"[DRY-RUN 开始] 计划迁移 {len(tables)} 张表 → {dst_driver}:{dst_db}")
    print("=" * 62)
    with src_engine.connect() as src:
        for t in tables:
            n = src.execute(text(f"SELECT COUNT(*) FROM {t.name}")).scalar()
            print(f"  - {t.name}: {n} 行")
    print(f"[DRY-RUN 完成] {len(tables)} 张表核对完毕（耗时 {_fmt_duration(time.perf_counter() - t_start)}），未触碰目标库")
    return 0


def run(args) -> int:
    """执行一次迁移。返回进程退出码（0 成功 / 1 校验失败）"""
    # 加载 Django 配置
    django_config = load_django_config()

    # 解析源库配置
    if args.src_driver:
        src_driver = _normalize_driver(args.src_driver)
        src_host = args.src_host or '127.0.0.1'
        src_port = _resolve_port(src_driver, args.src_port)
        src_db = args.src_db
        src_user = args.src_user or 'root'
        src_pass = args.src_pass or ''
    elif django_config:
        src_driver = django_engine_to_driver(django_config['engine'])
        src_host = django_config['host']
        src_port = int(django_config['port'])
        src_db = django_config['name']
        src_user = django_config['user']
        src_pass = django_config['password']
    else:
        print("[错误] 无法加载源库配置，请通过 --src-* 参数指定或确保 Django 配置可用")
        return 1

    # 解析目标库配置
    dst_driver = _normalize_driver(args.dst_driver)
    dst_host = args.dst_host
    dst_port = _resolve_port(dst_driver, args.dst_port)
    dst_db = args.dst_db
    dst_user = args.dst_user
    dst_pass = args.dst_pass

    # 获取表前缀
    prefix = args.prefix
    if prefix is None:
        prefix = django_config['prefix'] if django_config else 'django_'

    # 创建引擎
    src_engine = make_engine(src_driver, src_host, src_port, src_db, src_user, src_pass)
    dst_engine = make_engine(dst_driver, dst_host, dst_port, dst_db, dst_user, dst_pass)

    # 加载 Django 模型并创建 SQLAlchemy ORM 模型
    print("  · 加载 Django 模型定义...")
    load_django_models(prefix)
    print(f"  · 已加载 {len(Base.metadata.tables)} 张表的模型定义")

    # 获取迁移表列表
    tables = ordered_tables(prefix)
    if args.only_tables:
        wanted = {x.strip() for x in args.only_tables.split(",") if x.strip()}
        tables = [t for t in tables if t.name in wanted or t.name[len(prefix):] in wanted]

    try:
        src_insp = inspect(src_engine)
        present, missing = [], []
        for t in tables:
            if src_insp.has_table(t.name):
                present.append(t)
            else:
                missing.append(t)
        for t in missing:
            print(f"[SKIP] 源库不存在表 {t.name}，跳过")

        t_start = time.perf_counter()
        if args.dry_run:
            return _dry_run(src_engine, present, dst_driver, dst_db, t_start)

        print("=" * 62)
        print(f"[迁移开始] 源: {src_driver}:{src_db} → 目标: {dst_driver}:{dst_db}")
        note = "" if not missing else f"，源端缺失 {len(missing)} 张已跳过"
        print(f"[迁移开始] 待迁移 {len(present)} 张表{note}")
        print("=" * 62)

        # 1) 建目标库
        t0 = time.perf_counter()
        _phase_start(1, 5, "建目标库")
        create_database_if_missing(
            dst_driver, dst_host, dst_port,
            dst_db, dst_user, dst_pass,
            drop_first=args.drop_target_first,
            tables=tuple(t.name for t in present),
        )
        _phase_done(1, 5, "建目标库", t0)

        # 2) 建结构（使用 ORM 模型的 metadata，方言可移植）
        t0 = time.perf_counter()
        _phase_start(2, 5, f"建表结构（{len(present)} 张表 create_all）")
        Base.metadata.create_all(dst_engine)
        _phase_done(2, 5, f"建表结构（{len(present)} 张表 create_all）", t0, extra=f"，{len(present)} 张表")

        # 2.5) 清空目标表数据（如果指定了 --truncate-target）
        if args.truncate_target:
            print("  · 正在清空目标表数据...")
            with dst_engine.begin() as dst:
                # 按反向拓扑序清空（先清依赖表，再清基础表）
                for t in reversed(present):
                    try:
                        dst.execute(text(f"TRUNCATE TABLE {t.name} CASCADE"))
                        print(f"  · 已清空 {t.name}")
                    except Exception as e:
                        # 某些表可能不存在，忽略错误
                        print(f"  · [WARN] 清空 {t.name} 失败: {e}")

        # 3) 逐表迁移数据
        t0 = time.perf_counter()
        _phase_start(3, 5, "逐表迁移数据")
        total_rows = 0
        n_present = len(present)

        if args.single_transaction:
            with dst_engine.begin() as dst, src_engine.connect() as src:
                for idx, t in enumerate(present, 1):
                    if args.skip_existing and _existing_equal(src, dst_engine, t):
                        print(f"  · [{idx}/{n_present}] 跳过 {t.name}（源=目标）")
                        continue
                    print(f"  · [{idx}/{n_present}] 正在迁移 {t.name}...")
                    total_rows += pump_table(src_engine, dst_engine, t, args.chunk_size, dst_driver)
        else:
            for idx, t in enumerate(present, 1):
                if args.skip_existing:
                    with src_engine.connect() as src:
                        if _existing_equal(src, dst_engine, t):
                            print(f"  · [{idx}/{n_present}] 跳过 {t.name}（源=目标）")
                            continue
                print(f"  · [{idx}/{n_present}] 正在迁移 {t.name}...")
                total_rows += pump_table(src_engine, dst_engine, t, args.chunk_size, dst_driver)

        _phase_done(3, 5, "逐表迁移数据", t0, extra=f"，共 {total_rows} 行")

        # 4) 重置自增序列
        reset_errors = []
        t0 = time.perf_counter()
        _phase_start(4, 5, "重置自增序列")
        if dst_driver in ("postgresql", "mssql", "oracle"):
            # 每个表使用独立的事务，避免一个失败导致全部失败
            for t in present:
                try:
                    with dst_engine.begin() as dst:
                        reset_sequence(dst_driver, dst, t)
                except Exception as e:
                    reset_errors.append((t.name, str(e)))
                    print(f"  · [WARN] 序列重置失败 {t.name}: {e}")
            tail = f"，{len(reset_errors)} 张失败" if reset_errors else ""
            _phase_done(4, 5, "重置自增序列", t0, extra=tail)
        else:
            print("  · MySQL/SQLite 自动跟随，无需处理")
            _phase_done(4, 5, "重置自增序列", t0, extra="（无需处理）")

        # 5) 校验
        t0 = time.perf_counter()
        _phase_start(5, 5, "逐表行数校验")
        if not args.skip_verify:
            with src_engine.connect() as src, dst_engine.connect() as dst:
                failed = verify_counts(src, dst, present)
            if failed:
                _phase_done(5, 5, "逐表行数校验", t0, extra=f"，FAIL {len(failed)} 张")
                for name, s, d in failed:
                    print(f"  - {name}: source={s} target={d}")
                return 1
            _phase_done(5, 5, "逐表行数校验", t0, extra=f"，{len(present)} 张表全部 PASS")
        else:
            print("  · 已跳过（--skip-verify）")
            _phase_done(5, 5, "逐表行数校验", t0, extra="（已跳过）")

        print("=" * 62)
        print(f"[迁移完成] 成功：{total_rows} 行 / {len(present)} 张表 / 总耗时 {_fmt_duration(time.perf_counter() - t_start)}")
        print("=" * 62)
        return 0
    finally:
        src_engine.dispose()
        dst_engine.dispose()


def main():
    args = build_parser().parse_args()
    try:
        sys.exit(run(args))
    except KeyboardInterrupt:
        sys.exit(130)


if __name__ == "__main__":
    main()