#!/usr/bin/env python3
"""CMA 司库（资金）阶段一 —— 迁移脚本（阶段 1 落地版）

落地路径：/root/cma-management/backend/scripts/treasury_migrate.py
评审稿     ：wecom-fullstack/pending-tasks/_reports/treasury-stage1-migration-draft-20260915.py（md5 ceaa3b3bf563b8d664ddd7224f01a5bf，归档不回改）
G1 修正项  ：P0-1 up() 语句过滤器 / P0-2 计数口径 / P1-3 down() 部分态 / P1-4 down() dump 凭据 /
             P1-5 种子展开 / P2-6 data_status fail-safe（DDL 侧）/ P2-7 import_logs 指纹行 / P2-8 金额纪律

用法:
    python3 treasury_migrate.py --parse-only  # 零 DB：只解析 DDL，打印命中语句与期望计数（P0-1 复算证据）
    python3 treasury_migrate.py --check       # 只读预检：缺表/缺列/缺索引/种子行数（零写入，默认动作）
    python3 treasury_migrate.py --apply       # 建表 + 加列 + 索引 + 种子（幂等）
    python3 treasury_migrate.py --down        # 回滚（先 dump 后删；有数据时要求 --yes）

幂等契约（P0-2 计数口径 = 新增对象数，不是执行语句数）:
    首跑 --apply 期望：created_new=6 / added_cols=4 / added_indexes=1 / rules=8 行 / limits=6 行
    二跑 --apply 期望：created_new=0 / added_cols=0 / added_indexes=0 / 种子行数不变
    指纹行：import_logs(batch=treasury-ddl-<sha1[:10]>)，按 (import_type,batch) 先查后写，
            不依赖表上是否有唯一键（P2-7：entity_id=1 为约定——迁移指纹属实体无关元数据）
边界：只动 6 张 treasury_* 新表 + cash_plans 加列/加索引；不碰 cash.py / 既有端点契约 / 其他模块表结构。
"""
import argparse
import hashlib
import json
import os
import subprocess
import sys
import tempfile
from datetime import datetime

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

from sqlalchemy import create_engine, text  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
BACKUP_DIR = os.path.join(HERE, "backups")
DDL_FILE = os.path.join(HERE, "treasury_ddl.sql")

NEW_TABLES = [
    "treasury_accounts",
    "treasury_balances",
    "treasury_control_rules",
    "treasury_payments",
    "treasury_control_hits",
    "treasury_limits",
]
CASH_PLAN_NEW_COLS = ["account_id", "tier", "pool_id", "is_locked"]
CASH_PLAN_NEW_INDEX = "ix_cashplans_entity_acct"
EXPECTED_RULE_ROWS = 8      # e1 4 + e2 4
EXPECTED_LIMIT_ROWS = 6     # e1 3 + e2 3
CASH_PLAN_ALTER = """ALTER TABLE cash_plans
  ADD COLUMN account_id INT NULL COMMENT '资金来源/去向账户 treasury_accounts.id（软链，无物理外键）',
  ADD COLUMN tier      VARCHAR(20) NULL COMMENT 'hq/platform/grassroots（三位一体视角；可由 org_nodes.level 派生）',
  ADD COLUMN pool_id   INT NULL COMMENT '归属资金池（阶段三启用，本阶段恒 NULL）',
  ADD COLUMN is_locked TINYINT DEFAULT 0 COMMENT '是否已被预算/内控锁定'"""

# ─────────────────────────────────────────────────────────────
# 特权迁移连接（2026-09-16 并入）：与应用运行账号 cma_api 剥离，走有 DDL 权限的通道
# ─────────────────────────────────────────────────────────────
DEFAULT_SOCKET = "/var/run/mysqld/mysqld.sock"
FORBIDDEN_RUN_ACCOUNTS = {"cma_api"}
GRANT_APP_ACCOUNT = "cma_api"


def get_engine():
    """特权引擎：CMA_MIGRATE_DB_USER/PASS（TCP）优先，否则本机 root unix socket（无口令）。
    绝不使用应用运行账号（无 DDL）——与 db_migrate.py 同口径。"""
    from app import database as db
    db_name = (os.getenv("CMA_MIGRATE_DB_NAME") or getattr(db, "DB_NAME", None) or "cma")
    user = os.getenv("CMA_MIGRATE_DB_USER", "").strip()
    if user:
        if user in FORBIDDEN_RUN_ACCOUNTS:
            raise SystemExit(
                "[FAIL] CMA_MIGRATE_DB_USER=cma_api 是应用运行账号（无 DDL 权限），已拒绝执行")
        pw = os.getenv("CMA_MIGRATE_DB_PASS", "")
        host = os.getenv("CMA_MIGRATE_DB_HOST", "127.0.0.1")
        port = int(os.getenv("CMA_MIGRATE_DB_PORT", "3306"))
        return create_engine(
            f"mysql+pymysql://{user}:{pw}@{host}:{port}/{db_name}?charset=utf8mb4",
            pool_pre_ping=True)
    sock = os.getenv("CMA_MIGRATE_DB_SOCKET", DEFAULT_SOCKET)
    if not os.path.exists(sock):
        raise SystemExit(
            f"[FAIL] 未设 CMA_MIGRATE_DB_USER 且 unix socket 不存在：{sock}")
    return create_engine(
        f"mysql+pymysql://root@localhost/{db_name}?unix_socket={sock}&charset=utf8mb4",
        pool_pre_ping=True)


def grant_app_account(conn):
    """授权应用运行账号 cma_api 对 6 张 treasury 新表的 DML（幂等 GRANT，双 host）。
    必须由特权连接执行；否则应用端点读不到/写不进新表（1142/1044）。"""
    from app import database as db
    dbn = (os.getenv("CMA_MIGRATE_DB_NAME") or getattr(db, "DB_NAME", None) or "cma")
    granted = []
    for host in ("localhost", "127.0.0.1"):
        for t in NEW_TABLES:
            conn.execute(text(
                f"GRANT SELECT, INSERT, UPDATE, DELETE ON `{dbn}`.`{t}` "
                f"TO '{GRANT_APP_ACCOUNT}'@{host}"))
            granted.append(f"{t}@{host}")
    return granted


# ─────────────────────────────────────────────────────────────
# DDL 解析（P0-1：先剔除注释行再判定语句类型，不再用 startswith 裸判）
# ─────────────────────────────────────────────────────────────

def _strip_comment_lines(sql_text: str) -> str:
    """剔除整行注释（-- 开头）+ 行尾注释，避免注释行带动导致语句被跳过/误判。"""
    out = []
    for raw in sql_text.splitlines():
        line = raw
        idx = line.find("--")
        if idx >= 0:
            line = line[:idx]
        if not line.strip():
            continue
        out.append(line)
    return "\n".join(out)


def _split_statements(sql_text: str) -> list:
    return [s.strip() for s in _strip_comment_lines(sql_text).split(";") if s.strip()]


def _stmt_table_name(stmt: str, keyword: str):
    """取 `KEYWORD TABLE [IF NOT EXISTS] <name>` 中的表名。"""
    toks = stmt.replace("(", " ").replace("\n", " ").split()
    up = [t.upper() for t in toks]
    try:
        i = up.index(keyword)
    except ValueError:
        return None
    j = i + 1
    while j < len(up) and up[j] in ("IF", "NOT", "EXISTS"):
        j += 1
    return toks[j].strip("`") if j < len(toks) else None


def parse_ddl(path=DDL_FILE):
    """解析 DDL → {creates: [(table, stmt)], seeds: [stmt], alter: str|None, index: str|None}"""
    sql = open(path, encoding="utf-8").read()
    creates, seeds, alter, index_stmt = [], [], None, None
    for stmt in _split_statements(sql):
        head = stmt.upper()
        if head.startswith("CREATE TABLE"):
            creates.append((_stmt_table_name(stmt, "TABLE"), stmt))
        elif head.startswith("INSERT INTO"):
            seeds.append(stmt)
        elif head.startswith("ALTER TABLE CASH_PLANS ADD KEY"):
            index_stmt = stmt
        elif head.startswith("ALTER TABLE CASH_PLANS"):
            alter = stmt
    return {"creates": creates, "seeds": seeds, "alter": alter, "index": index_stmt,
            "ddl_text": sql}


def assert_ddl_parseable(p):
    """P0-1 断言：6 条 CREATE 全命中且表名与 NEW_TABLES 一致；种子 2 条；ALTER/B1 各 1 条。"""
    names = [t for t, _ in p["creates"]]
    errs = []
    if len(p["creates"]) != len(NEW_TABLES):
        errs.append(f"CREATE TABLE 命中数 {len(p['creates'])} != {len(NEW_TABLES)}（{names}）")
    if set(names) != set(NEW_TABLES):
        errs.append(f"CREATE TABLE 表名集合不一致: {sorted(names)} vs {sorted(NEW_TABLES)}")
    if len(p["seeds"]) != 2:
        errs.append(f"种子 INSERT 语句数 {len(p['seeds'])} != 2")
    if not p["alter"]:
        errs.append("DDL §B 未解析到 ALTER TABLE cash_plans 加列语句")
    if not p["index"]:
        errs.append("DDL §B1 未解析到 ALTER TABLE cash_plans ADD KEY 语句")
    if p["alter"] and p["alter"].strip() not in CASH_PLAN_ALTER.strip():
        errs.append("DDL §B ALTER 文本与脚本 CASH_PLAN_ALTER 常量漂移（指纹口径会不一致）")
    return errs


def _fingerprint(texts):
    canon = json.dumps(sorted(texts), ensure_ascii=False, separators=(",", ":"))
    return hashlib.sha1(canon.encode("utf-8")).hexdigest()[:10]


# ─────────────────────────────────────────────────────────────
# information_schema 内省（只读）
# ─────────────────────────────────────────────────────────────

def _introspect(conn):
    rows = conn.execute(text(
        "SELECT TABLE_NAME FROM information_schema.TABLES "
        "WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME LIKE 'treasury%'"
    )).fetchall()
    tables = {r[0] for r in rows}
    cols = conn.execute(text(
        "SELECT COLUMN_NAME FROM information_schema.COLUMNS "
        "WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='cash_plans'"
    )).fetchall()
    colset = {r[0] for r in cols}
    idx = conn.execute(text(
        "SELECT DISTINCT INDEX_NAME FROM information_schema.STATISTICS "
        "WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='cash_plans'"
    )).fetchall()
    return tables, colset, {r[0] for r in idx}


def _seed_counts(conn):
    counts = {}
    for tbl, key in (("treasury_control_rules", "rules_rows"), ("treasury_limits", "limits_rows")):
        try:
            counts[key] = conn.execute(text(f"SELECT COUNT(*) FROM {tbl}")).scalar()
        except Exception:
            counts[key] = None
    return counts


def check(conn):
    """只读预检 —— 缺表 / 缺列 / 缺索引 / 存量行数 / 种子行数；零写入。"""
    tables, cols, idx = _introspect(conn)
    cash_rows = conn.execute(text("SELECT COUNT(*) FROM cash_plans")).scalar()
    report = {
        "missing_tables": [t for t in NEW_TABLES if t not in tables],
        "existing_tables": sorted(tables),
        "missing_cash_plan_cols": [c for c in CASH_PLAN_NEW_COLS if c not in cols],
        "missing_cash_plan_index": [] if CASH_PLAN_NEW_INDEX in idx else [CASH_PLAN_NEW_INDEX],
        "cash_plans_rows": cash_rows,
    }
    report.update(_seed_counts(conn))
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return report


# ─────────────────────────────────────────────────────────────
# UP / DOWN
# ─────────────────────────────────────────────────────────────

def _write_fingerprint(conn, batch, tables_created):
    """P2-7：指纹行按 (import_type,batch) 先查后写（不依赖唯一键；entity_id=1 为约定）。"""
    cols = {r[0] for r in conn.execute(text(
        "SELECT COLUMN_NAME FROM information_schema.COLUMNS "
        "WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='import_logs'"
    )).fetchall()}
    need = {"entity_id", "filename", "batch", "total_rows", "success_rows", "import_type"}
    if not need <= cols:
        raise RuntimeError(f"import_logs 结构不符，缺列 {sorted(need - cols)}（fail-closed，不静默跳过）")
    row = conn.execute(text(
        "SELECT id FROM import_logs WHERE import_type='treasury-ddl' AND batch=:b ORDER BY id LIMIT 1"
    ), {"b": batch}).fetchone()
    if row:
        conn.execute(text(
            "UPDATE import_logs SET total_rows=:t, success_rows=:s WHERE id=:i"
        ), {"t": tables_created, "s": tables_created, "i": row[0]})
        return "updated"
    conn.execute(text(
        "INSERT INTO import_logs (entity_id, filename, batch, total_rows, success_rows, import_type, created_by) "
        "VALUES (1, 'treasury_ddl.sql', :b, :t, :s, 'treasury-ddl', 'treasury_migrate')"
    ), {"b": batch, "t": tables_created, "s": tables_created})
    return "inserted"


def up(conn):
    p = parse_ddl()
    errs = assert_ddl_parseable(p)
    if errs:
        for e in errs:
            print("❌ DDL 解析失败:", e, file=sys.stderr)
        return 3

    before_tables, before_cols, before_idx = _introspect(conn)

    # ① 建表（6 条 CREATE，IF NOT EXISTS 幂等）
    for _table, stmt in p["creates"]:
        conn.execute(text(stmt))

    # ② 条件加列（列级幂等）
    missing_cols = [c for c in CASH_PLAN_NEW_COLS if c not in before_cols]
    if missing_cols:
        conn.execute(text(p["alter"]))

    # ③ 条件加索引（D-7）
    if CASH_PLAN_NEW_INDEX not in before_idx:
        conn.execute(text(p["index"]))

    # ③b 授权应用账号 cma_api 对本 6 表 DML（幂等 GRANT，须特权连接）
    granted = grant_app_account(conn)

    # ④ 种子（ON DUPLICATE KEY UPDATE，数据级幂等）
    for stmt in p["seeds"]:
        conn.execute(text(stmt))

    conn.commit()

    after_tables, after_cols, after_idx = _introspect(conn)
    created_new = len(after_tables - before_tables)
    added_cols = len(after_cols - before_cols)
    added_indexes = len(after_idx - before_idx)
    counts = _seed_counts(conn)

    # ⑤ 自检断言：缺表即非零退出（P0-1，禁静默成功）
    post = check(conn)
    if post["missing_tables"]:
        print(f"❌ up() 后仍有缺表 {post['missing_tables']}（非零退出）", file=sys.stderr)
        return 4

    # ⑥ 指纹行
    batch = "treasury-ddl-" + _fingerprint([p["ddl_text"], CASH_PLAN_ALTER])
    action = _write_fingerprint(conn, batch, created_new or len(NEW_TABLES))
    conn.commit()

    print(json.dumps({
        "phase": "up",
        "created_new": created_new,
        "added_cols": added_cols,
        "added_indexes": added_indexes,
        "seeded_statements": len(p["seeds"]),
        "rules_rows": counts["rules_rows"],
        "limits_rows": counts["limits_rows"],
        "import_logs_row": action,
        "batch": batch,
        "granted": granted,
        "missing_tables": post["missing_tables"],
        "cash_plans_rows": post["cash_plans_rows"],
    }, ensure_ascii=False))
    return 0


def _dump_cash_plans(out_path):
    """P1-4：写路径账号 + 临时 defaults-file 注入口令（不硬编码、不回显）；失败即抛异常（fail-closed）。"""
    from app import database as db

    fd, cnf = tempfile.mkstemp(prefix=".cma-dump-", suffix=".cnf")
    try:
        with os.fdopen(fd, "w") as f:
            f.write("[client]\nuser=%s\npassword=%s\nhost=%s\nport=%s\n"
                    % (db.DB_USER, db.DB_PASS, db.DB_HOST, db.DB_PORT))
        os.chmod(cnf, 0o600)
        with open(out_path, "w") as fh:
            proc = subprocess.run(
                ["mysqldump", f"--defaults-extra-file={cnf}", "--single-transaction",
                 "--skip-comments", db.DB_NAME, "cash_plans"],
                stdout=fh, stderr=subprocess.PIPE, text=True)
        if proc.returncode != 0:
            raise RuntimeError("mysqldump 失败(rc=%s): %s" % (proc.returncode, (proc.stderr or "").strip()[:300]))
    finally:
        try:
            os.remove(cnf)
        except OSError:
            pass


def down(conn):
    """回滚：先 dump（fail-closed）→ 逐列删扩展列（P1-3）→ 删索引 → 删表（IF EXISTS）。"""
    os.makedirs(BACKUP_DIR, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    dump = os.path.join(BACKUP_DIR, f"cash_plans-before-rollback-{stamp}.sql")
    try:
        _dump_cash_plans(dump)
    except Exception as exc:  # fail-closed：dump 不成功就不允许删列
        print(f"❌ 备份失败，回滚中止（未做任何结构变更）: {exc}", file=sys.stderr)
        return 5
    print("backup:", dump, os.path.getsize(dump), "bytes")

    tables, cols, idx = _introspect(conn)
    dropped = []
    for col in CASH_PLAN_NEW_COLS:            # P1-3：按列存在性逐列 DROP（部分态不再 1091）
        if col in cols:
            conn.execute(text(f"ALTER TABLE cash_plans DROP COLUMN {col}"))
            dropped.append(col)
    if CASH_PLAN_NEW_INDEX in idx:
        conn.execute(text(f"ALTER TABLE cash_plans DROP INDEX {CASH_PLAN_NEW_INDEX}"))
        dropped.append(CASH_PLAN_NEW_INDEX)
    for t in reversed(NEW_TABLES):
        conn.execute(text(f"DROP TABLE IF EXISTS {t}"))
    conn.commit()
    print(json.dumps({"phase": "down", "dropped_cash_plan_objects": dropped,
                      "dropped_tables": list(reversed(NEW_TABLES)), "backup": dump},
                     ensure_ascii=False))
    return 0


# ─────────────────────────────────────────────────────────────
# 入口
# ─────────────────────────────────────────────────────────────

def parse_only():
    """零 DB 的 DDL 解析复算（P0-1 证据）：命中 6 条 CREATE + 表名清单 + 种子语句。"""
    p = parse_ddl()
    errs = assert_ddl_parseable(p)
    out = {
        "ddl_file": DDL_FILE,
        "ddl_md5": hashlib.md5(open(DDL_FILE, "rb").read()).hexdigest(),
        "statements_total": len(_split_statements(p["ddl_text"])),
        "create_table_hits": len(p["creates"]),
        "create_table_names": [t for t, _ in p["creates"]],
        "seed_statements": len(p["seeds"]),
        "has_cash_plan_alter": bool(p["alter"]),
        "has_cash_plan_index": bool(p["index"]),
        "expected_first_run": {"created_new": 6, "added_cols": 4, "added_indexes": 1,
                               "rules_rows": EXPECTED_RULE_ROWS, "limits_rows": EXPECTED_LIMIT_ROWS},
        "expected_second_run": {"created_new": 0, "added_cols": 0, "added_indexes": 0,
                                "rules_rows": EXPECTED_RULE_ROWS, "limits_rows": EXPECTED_LIMIT_ROWS},
        "errors": errs,
    }
    print(json.dumps(out, ensure_ascii=False, indent=2))
    return 1 if errs else 0


def main():
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group()
    g.add_argument("--parse-only", action="store_true", help="零 DB：只解析 DDL（P0-1 复算证据）")
    g.add_argument("--check", action="store_true", help="只读预检（默认）")
    g.add_argument("--apply", action="store_true")
    g.add_argument("--down", action="store_true")
    ap.add_argument("--yes", action="store_true", help="回滚确认（表内有数据时必需）")
    args = ap.parse_args()

    if args.parse_only:
        return parse_only()

    engine = get_engine()
    with engine.connect() as conn:
        if args.apply:
            return up(conn)
        if args.down:
            tables, _c, _i = _introspect(conn)
            payload = 0
            if {"treasury_payments", "treasury_balances"} <= tables:
                payload = conn.execute(text(
                    "SELECT (SELECT COUNT(*) FROM treasury_payments) + "
                    "(SELECT COUNT(*) FROM treasury_balances)")).scalar()
            if payload and not args.yes:
                print(f"REFUSE: 已有 {payload} 行业务数据，回滚需显式 --yes", file=sys.stderr)
                return 2
            return down(conn)
        return 0 if check(conn) is not None else 1


if __name__ == "__main__":
    sys.exit(main())