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

功能：
1. 创建数据库（默认 utf8mb4）
2. 导入 document/djangoadmin.fastapi.elevue.sql（若存在）
3. 用模型 metadata 补建 SQL 中缺失的表（如 fastapi_login_log / fastapi_operation_log）

用法：
    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.dirname(os.path.dirname(os.path.abspath(__file__))))


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 main():
    parser = argparse.ArgumentParser(description="初始化数据库")
    parser.add_argument("--drop-first", action="store_true", help="先删除已存在库再重建")
    args = parser.parse_args()

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

    import pymysql
    from sqlalchemy import create_engine
    from core.database import Base

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

    create_app()

    root_conn = pymysql.connect(host=DB_HOST, port=int(DB_PORT), user=DB_USERNAME, password=DB_PASSWORD)
    cur = root_conn.cursor()
    if args.drop_first:
        cur.execute(f"DROP DATABASE IF EXISTS `{DB_DATABASE}`")
    cur.execute(
        f"CREATE DATABASE IF NOT EXISTS `{DB_DATABASE}` "
        f"DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci"
    )
    cur.execute(f"USE `{DB_DATABASE}`")
    print(f"database {DB_DATABASE} ready (prefix={DB_PREFIX})")

    sql_path = os.path.join("document", "djangoadmin.fastapi.elevue.sql")
    if os.path.exists(sql_path):
        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"imported {done} statements, errors={errors}")
    else:
        print(f"!! 未找到 {sql_path}，跳过 SQL 导入")

    # 用模型补建缺失表
    engine = create_engine(
        f"mysql+pymysql://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_DATABASE}?charset=utf8mb4"
    )
    Base.metadata.create_all(engine)
    from sqlalchemy import inspect

    actual = set(inspect(engine).get_table_names())
    missing = set(Base.metadata.tables) - actual
    print(f"tables in DB: {len(actual)}; still missing: {sorted(missing)}")
    engine.dispose()
    cur.close()
    root_conn.close()


if __name__ == "__main__":
    main()