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

# +======================================================================
# | 模块: 数据库初始化脚本
# | 说明: 按 DB_DRIVER 分方言创建数据库并建表
# +======================================================================

"""
数据库初始化脚本
=============================================================

按 DB_DRIVER 分方言创建数据库并建表：
- mysql:      创建库（utf8mb4）→ 导入 document/mysql/djangoadmin.fastapi.elevue.sql（若存在）
              → 用模型 metadata 补建缺失的表
- postgresql: 连 postgres 维护库，无则建库 → 导入 document/postgresql/ 下的 SQL（若存在）→ create_all
- mssql:      连 master 库，无则建库 → 导入 document/sqlserver/ 下的 SQL（若存在）→ create_all
- oracle:     连 service_name 校验连通（Oracle 无"建库"概念，schema 即连接用户，
              需 DBA 预建）→ create_all 建表 → 自动补 sequence + BEFORE INSERT
              触发器（Oracle create_all 不生成 Integer 主键自增）
- sqlite:     直接 create_all 生成数据库文件（DB_DATABASE 即数据库文件完整路径）

SQL 脚本按数据库类型分目录存放在 document/ 下：
    document/mysql/djangoadmin.fastapi.elevue.sql
    document/postgresql/djangoadmin.fastapi.elevue.sql
    document/sqlserver/djangoadmin.fastapi.elevue.sql
    document/sqlite/djangoadmin.fastapi.elevue.sql
    document/oracle/djangoadmin.fastapi.elevue.sql

用法：
    python scripts/init_db.py                     # 使用 .env 中的数据库配置
    python scripts/init_db.py --drop-first         # 先删除已存在的同名库再重建
"""

import argparse
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))

# Windows 控制台默认 GBK(cp936)：把 stdout/stderr 切到 UTF-8，避免中文打印乱码/崩溃
if sys.stdout and hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')
    sys.stderr.reconfigure(encoding='utf-8')

from sqlalchemy import text


# ============================================================
# 辅助函数
# ============================================================
def _stamp_alembic():
    """建库完成后将库标记到 alembic 基线，后续增量走迁移。失败不影响建库"""
    try:
        from alembic.config import Config
        from alembic import command

        root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        cfg = Config(os.path.join(root, "alembic.ini"))
        command.stamp(cfg, "head")
        print("[INFO] Alembic 版本标记完成")
    except Exception as e:
        print(f"[WARN] Alembic 版本标记失败（可稍后手动执行 alembic stamp head）: {e}")


def split_statements(content: str):
    """按 ';\\n' 切分语句，并剥离 -- 注释行，避免注释+语句整块被跳过"""
    for block in content.split(";\n"):
        body = []
        for line in block.splitlines():
            s = line.strip()
            if s.startswith("--") or s.startswith("#"):
                continue
            body.append(line)
        stmt = "\n".join(body).strip()
        if stmt:
            yield stmt


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 create_all 后补自增机制：重建 sequence + BEFORE INSERT 触发器"""
    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 main():
    parser = argparse.ArgumentParser(description="初始化数据库")
    parser.add_argument("--drop-first", action="store_true", help="先删除已存在库再重建")
    args = parser.parse_args()

    # 加载 .env 数据库配置
    from core.config import (
        DB_DRIVER, DB_HOST, DB_PORT, DB_DATABASE, DB_USERNAME, DB_PASSWORD,
        DB_PREFIX, SQLALCHEMY_DATABASE_URL,
    )

    from sqlalchemy import create_engine, inspect
    from core.database import Base

    # 加载全部模型（导入 create_app 会连带导入所有模块）
    from core.app import create_app

    create_app()

    # ----------------------------------------------------------
    # 启动日志
    # ----------------------------------------------------------
    print("[INFO] ============================================================")
    print("[INFO] 数据库初始化开始")
    print("[INFO] ============================================================")
    print(f"[INFO] 检测数据库驱动: {DB_DRIVER}")

    # 驱动 → 可读名称
    driver_label = {
        'mysql': 'MySQL', 'postgresql': 'PostgreSQL',
        'mssql': 'SQL Server', 'oracle': 'Oracle', 'sqlite': 'SQLite',
    }

    if DB_DRIVER != 'sqlite':
        print(f"[INFO] 连接 {driver_label.get(DB_DRIVER, DB_DRIVER)} 服务器: {DB_HOST}:{DB_PORT}")
    print(f"[INFO] 数据库配置: {DB_USERNAME}@{DB_HOST}:{DB_PORT}/{DB_DATABASE}")
    print(f"[INFO] 数据表前缀: {DB_PREFIX}")

    # SQL 脚本路径：document/{driver}/djangoadmin.fastapi.elevue.sql
    sql_filename = "djangoadmin.fastapi.elevue.sql"
    driver_dir_map = {
        'mysql': 'mysql',
        'postgresql': 'postgresql',
        'mssql': 'sqlserver',
        'oracle': 'oracle',
        'sqlite': 'sqlite',
    }
    sql_path = os.path.join("document", driver_dir_map.get(DB_DRIVER, DB_DRIVER), sql_filename)

    # ----------------------------------------------------------
    # 按驱动建库
    # ----------------------------------------------------------
    if DB_DRIVER == 'mysql':
        import pymysql

        root_conn = pymysql.connect(host=DB_HOST, port=int(DB_PORT), user=DB_USERNAME, password=DB_PASSWORD)
        cur = root_conn.cursor()

        # 检查数据库是否已存在
        cur.execute(f"SHOW DATABASES LIKE '{DB_DATABASE}'")
        db_exists = cur.fetchone() is not None

        if args.drop_first and db_exists:
            print(f"[INFO] 删除已有数据库: {DB_DATABASE}")
            cur.execute(f"DROP DATABASE IF EXISTS `{DB_DATABASE}`")
            db_exists = False

        if db_exists:
            print(f"[INFO] 数据库 {DB_DATABASE} 已存在，跳过创建")
        else:
            cur.execute(
                f"CREATE DATABASE IF NOT EXISTS `{DB_DATABASE}` "
                f"DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci"
            )
            print(f"[INFO] 创建数据库: {DB_DATABASE}")

        cur.execute(f"USE `{DB_DATABASE}`")

        # 导入 SQL 脚本
        if os.path.exists(sql_path):
            print(f"[INFO] 查找数据库脚本: {sql_path}")
            print("[INFO] 导入 SQL 脚本...")
            with open(sql_path, "r", encoding="utf-8") as f:
                content = f.read()
            done = errors = 0
            for stmt in split_statements(content):
                try:
                    cur.execute(stmt)
                    done += 1
                except Exception as e:
                    errors += 1
                    print(f"[SKIP] {e}\n  {stmt[:100]}")
            root_conn.commit()
            print(f"[INFO] SQL 脚本导入完成: {done} 条语句")
        else:
            print(f"[WARN] 未找到 SQL 脚本: {sql_path}，跳过导入")

        cur.close()
        root_conn.close()

    elif DB_DRIVER == 'postgresql':
        import psycopg

        # 连 postgres 维护库创建目标库；DB_DATABASE 含点号，标识符必须双引号
        conn = psycopg.connect(
            host=DB_HOST, port=int(DB_PORT), user=DB_USERNAME,
            password=DB_PASSWORD, dbname="postgres", autocommit=True,
        )
        with conn.cursor() as cur:
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (DB_DATABASE,))
            exists = cur.fetchone() is not None

            if args.drop_first and exists:
                print(f"[INFO] 删除已有数据库: {DB_DATABASE}")
                cur.execute(f'DROP DATABASE "{DB_DATABASE}" WITH (FORCE)')
                exists = False

            if exists:
                print(f"[INFO] 数据库 {DB_DATABASE} 已存在，跳过创建")
            else:
                cur.execute(f'CREATE DATABASE "{DB_DATABASE}"')
                print(f"[INFO] 创建数据库: {DB_DATABASE}")
        conn.close()

        # 导入 PostgreSQL SQL 脚本
        if os.path.exists(sql_path):
            print(f"[INFO] 查找数据库脚本: {sql_path}")
            print("[INFO] 导入 SQL 脚本...")
            import psycopg as pg
            pg_conn = pg.connect(
                host=DB_HOST, port=int(DB_PORT), user=DB_USERNAME,
                password=DB_PASSWORD, dbname=DB_DATABASE, autocommit=True,
            )
            with pg_conn.cursor() as cur:
                with open(sql_path, "r", encoding="utf-8") as f:
                    cur.execute(f.read())
            pg_conn.close()
            print("[INFO] SQL 脚本导入完成")
        else:
            print(f"[WARN] 未找到 SQL 脚本: {sql_path}，跳过导入")

    elif DB_DRIVER == 'mssql':
        import pymssql

        # 连 master 维护库创建目标库；标识符用方括号引用
        conn = pymssql.connect(
            server=DB_HOST, user=DB_USERNAME, password=DB_PASSWORD,
            port=int(DB_PORT), database="master", autocommit=True,
        )
        cur = conn.cursor()

        # 检查数据库是否已存在
        cur.execute(f"SELECT DB_ID(N'{DB_DATABASE}')")
        db_exists = cur.fetchone()[0] is not None

        if args.drop_first and db_exists:
            print(f"[INFO] 删除已有数据库: {DB_DATABASE}")
            cur.execute(
                f"ALTER DATABASE [{DB_DATABASE}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE"
            )
            cur.execute(f"DROP DATABASE [{DB_DATABASE}]")
            db_exists = False

        if db_exists:
            print(f"[INFO] 数据库 {DB_DATABASE} 已存在，跳过创建")
        else:
            cur.execute(f"CREATE DATABASE [{DB_DATABASE}]")
            print(f"[INFO] 创建数据库: {DB_DATABASE}")
        conn.close()

        # 导入 SQL Server SQL 脚本
        if os.path.exists(sql_path):
            print(f"[INFO] 查找数据库脚本: {sql_path}")
            print("[INFO] 导入 SQL 脚本...")
            mssql_conn = pymssql.connect(
                server=DB_HOST, user=DB_USERNAME, password=DB_PASSWORD,
                port=int(DB_PORT), database=DB_DATABASE, autocommit=True,
            )
            mssql_cur = mssql_conn.cursor()
            with open(sql_path, "r", encoding="utf-8") as f:
                content = f.read()
            done = errors = 0
            for stmt in split_statements(content):
                try:
                    mssql_cur.execute(stmt)
                    done += 1
                except Exception as e:
                    errors += 1
                    print(f"[SKIP] {e}\n  {stmt[:100]}")
            mssql_conn.commit()
            mssql_conn.close()
            print(f"[INFO] SQL 脚本导入完成: {done} 条语句")
        else:
            print(f"[WARN] 未找到 SQL 脚本: {sql_path}，跳过导入")

    elif DB_DRIVER == 'oracle':
        # Oracle 的"库"即连接的服务名（service_name），schema 是连接用户，由 DBA 预建，
        # 无"建库"概念；此处仅校验连通性，建表走下方 create_all
        import oracledb

        conn = oracledb.connect(
            user=DB_USERNAME, password=DB_PASSWORD,
            dsn=f"{DB_HOST}:{DB_PORT}/{DB_DATABASE}",
        )
        conn.close()
        print(f"[INFO] Oracle 服务 {DB_DATABASE} 连通 (schema={DB_USERNAME})")

        # 导入 Oracle SQL 脚本
        if os.path.exists(sql_path):
            print(f"[INFO] 查找数据库脚本: {sql_path}")
            print("[INFO] 导入 SQL 脚本...")
            ora_conn = oracledb.connect(
                user=DB_USERNAME, password=DB_PASSWORD,
                dsn=f"{DB_HOST}:{DB_PORT}/{DB_DATABASE}",
            )
            ora_cur = ora_conn.cursor()
            with open(sql_path, "r", encoding="utf-8") as f:
                content = f.read()
            done = errors = 0
            for stmt in split_statements(content):
                try:
                    ora_cur.execute(stmt)
                    done += 1
                except Exception as e:
                    errors += 1
                    print(f"[SKIP] {e}\n  {stmt[:100]}")
            ora_conn.commit()
            ora_cur.close()
            ora_conn.close()
            print(f"[INFO] SQL 脚本导入完成: {done} 条语句")
        else:
            print(f"[WARN] 未找到 SQL 脚本: {sql_path}，跳过导入")

    else:  # sqlite（DB_DATABASE 即数据库文件完整路径）
        from pathlib import Path

        file = Path(DB_DATABASE).expanduser()
        print(f"[INFO] SQLite 数据库文件: {file}")

        if file.parent.as_posix() != '.':
            file.parent.mkdir(parents=True, exist_ok=True)
        if args.drop_first and file.exists():
            file.unlink()
            print(f"[INFO] 删除已有数据库文件: {file}")

        # 导入 SQLite SQL 脚本
        if os.path.exists(sql_path):
            print(f"[INFO] 查找数据库脚本: {sql_path}")
            print("[INFO] 导入 SQL 脚本...")
            import sqlite3 as sqlite_mod
            sq_conn = sqlite_mod.connect(str(file))
            with open(sql_path, "r", encoding="utf-8") as f:
                content = f.read()
            done = errors = 0
            for stmt in split_statements(content):
                try:
                    sq_conn.execute(stmt)
                    done += 1
                except Exception as e:
                    errors += 1
                    print(f"[SKIP] {e}\n  {stmt[:100]}")
            sq_conn.commit()
            sq_conn.close()
            print(f"[INFO] SQL 脚本导入完成: {done} 条语句")
        else:
            print(f"[WARN] 未找到 SQL 脚本: {sql_path}，跳过导入")

    # ----------------------------------------------------------
    # 用模型补建缺失表
    # ----------------------------------------------------------
    print("[INFO] 创建数据表...")

    engine = create_engine(SQLALCHEMY_DATABASE_URL)

    # 记录 create_all 前已存在的表，用于差量统计新建的表
    existing_before = set(inspect(engine).get_table_names())

    Base.metadata.create_all(engine)

    # 差量：打印新建的表
    existing_after = set(inspect(engine).get_table_names())
    newly_created = sorted(existing_after - existing_before)
    if newly_created:
        for t in newly_created:
            print(f"[INFO] 表 {t} 创建成功")
    else:
        print("[INFO] 所有数据表均已存在，无需创建")

    missing = set(Base.metadata.tables) - existing_after
    total_tables = len(existing_after)
    print(f"[INFO] 数据表创建完成: 共 {total_tables} 张表")
    if missing:
        print(f"[WARN] 以下表仍缺失: {sorted(missing)}")

    # Oracle：create_all 对 Integer 主键只生成 INTEGER NOT NULL（无自增），
    # 补 sequence + BEFORE INSERT 触发器
    if DB_DRIVER == 'oracle':
        print("[INFO] 补建 Oracle 自增 sequence + 触发器...")
        with engine.begin() as conn:
            for table_name in Base.metadata.tables:
                ensure_oracle_autoincrement(conn, table_name)
        print(f"[INFO] 已为 {len(Base.metadata.tables)} 张表补建自增 sequence + 触发器")

    engine.dispose()

    # ----------------------------------------------------------
    # Alembic 版本标记
    # ----------------------------------------------------------
    print("[INFO] 标记 Alembic 版本...")
    _stamp_alembic()

    # ----------------------------------------------------------
    # 完成
    # ----------------------------------------------------------
    print("[INFO] ============================================================")
    print("[INFO] 数据库初始化完成！")
    print("[INFO] ============================================================")


if __name__ == "__main__":
    main()