# -*- coding: utf-8 -*-
"""
建表脚本: 把 ddl_pms_v1.sql 应用到 153 代理侧
===============================================
运行:
  docker compose run --rm pms-web python scripts/init_db.py          # 演练, 只打印不执行
  docker compose run --rm pms-web python scripts/init_db.py --yes    # 实际建表
  docker compose run --rm pms-web python scripts/init_db.py --yes --only pms_command,pms_plan

说明:
  * 全部语句为 `CREATE TABLE IF NOT EXISTS`, 重复执行安全 (已存在的表不会被改动)。
  * 直接用引擎执行, 绕过 db.session 的单表守卫 —— DDL 本身就只涉及一张表,
    但语句里的中文注释可能撞上守卫的关键字启发式, 没必要让它误伤。
  * ShardingSphere-Proxy 对 DDL 的支持视配置而定。若某条语句被代理拒绝,
    脚本会把完整语句打出来, 直接拿去物理库 (my_quant_db) 执行即可, 其余语句不受影响。
"""
import argparse
import os
import re
import sys

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

DDL_FILE = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                        "ddl_pms_v1.sql")


def split_sql(text: str) -> list:
    """按分号切语句, 但**跳过字符串字面量里的分号**。

    DDL 的列注释里有 `COMMENT '参数命令: EFFECTIVE/SUPERSEDED; 任务命令: PENDING/...'`,
    直接 split(";") 会把一条 CREATE TABLE 劈成三段残句 —— 必须带引号状态机。
    """
    out, buf, in_str, i = [], [], False, 0
    while i < len(text):
        ch = text[i]
        if in_str:
            if ch == "\\" and i + 1 < len(text):        # 反斜杠转义
                buf.append(text[i:i + 2])
                i += 2
                continue
            if ch == "'":
                if i + 1 < len(text) and text[i + 1] == "'":   # '' 转义
                    buf.append("''")
                    i += 2
                    continue
                in_str = False
            buf.append(ch)
        elif ch == "'":
            in_str = True
            buf.append(ch)
        elif ch == ";":
            out.append("".join(buf))
            buf = []
        else:
            buf.append(ch)
        i += 1
    out.append("".join(buf))
    return out


def parse_statements(sql_text: str) -> list:
    """去掉整行 `--` 注释后切语句。返回 [(表名, 语句)]; 非 CREATE TABLE 的残句会被标 '?'。"""
    body = "\n".join(ln for ln in sql_text.splitlines() if not ln.strip().startswith("--"))
    out = []
    for stmt in split_sql(body):
        s = stmt.strip()
        if not s:
            continue
        m = re.search(r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?`?([A-Za-z0-9_]+)`?", s, re.I)
        out.append((m.group(1) if m else "?", s))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--yes", action="store_true", help="确认执行 (缺省只演练)")
    ap.add_argument("--only", default="", help="只处理指定表, 逗号分隔")
    args = ap.parse_args()

    if not os.path.exists(DDL_FILE):
        print(f"FAIL  找不到 DDL 文件: {DDL_FILE}")
        sys.exit(1)
    with open(DDL_FILE, encoding="utf-8") as f:
        stmts = parse_statements(f.read())
    only = {x.strip() for x in args.only.split(",") if x.strip()}
    if only:
        stmts = [(t, s) for t, s in stmts if t in only]
    if not stmts:
        print("FAIL  没有可执行的语句 (检查 --only 是否写对)")
        sys.exit(1)
    broken = [s for t, s in stmts if t == "?"]
    if broken:
        print(f"FAIL  解析出 {len(broken)} 条非 CREATE TABLE 残句, 说明 DDL 切分有误, "
              f"已中止 (不执行任何语句)。首条残句:\n{broken[0][:200]}")
        sys.exit(1)

    print(f"DDL 文件: {DDL_FILE}")
    print(f"待建表 {len(stmts)} 张: {', '.join(t for t, _ in stmts)}")
    if not args.yes:
        print("\n[演练模式] 未执行任何语句。确认无误后加 --yes 重跑:")
        print("  docker compose run --rm pms-web python scripts/init_db.py --yes")
        return

    # 连库依赖放到演练之后再导入: 演练模式只做语句切分校验, 无依赖环境也能跑
    from sqlalchemy import text
    from app.db.session import get_engine

    eng = get_engine("proxy")
    okc, failed = 0, []
    print()
    for tbl, stmt in stmts:
        try:
            with eng.begin() as c:
                c.execute(text(stmt))
            print(f"  OK    {tbl}")
            okc += 1
        except Exception as e:
            print(f"  FAIL  {tbl}: {type(e).__name__}: {e}")
            failed.append((tbl, stmt, str(e)))

    print("\n[验证] 逐表 COUNT(*)")
    missing = []
    for tbl, _ in stmts:
        try:
            with eng.connect() as c:
                n = c.execute(text(f"SELECT COUNT(*) FROM {tbl}")).scalar()
            print(f"  OK    {tbl:<20} {n} 行")
        except Exception as e:
            print(f"  FAIL  {tbl:<20} {type(e).__name__}: {e}")
            missing.append(tbl)

    print("\n" + "-" * 62)
    if failed:
        print(f"以下 {len(failed)} 张表建失败, 完整语句如下 —— 可直接拿到物理库执行:")
        for tbl, stmt, err in failed:
            print(f"\n### {tbl}  ({err.splitlines()[0]})\n{stmt};")
    if missing:
        print(f"\nFAILED: {len(missing)} 张表仍不可用: {', '.join(missing)}")
        sys.exit(1)
    print(f"ALL OK: {okc} 张表就绪, 接着跑 python scripts/check_db.py 做完整自检")


if __name__ == "__main__":
    main()