"""波次5.4：业务必填列 NOT NULL 对齐（幂等，仅 PostgreSQL）。

背景：ORM 声明 NOT NULL、线上库却允许 NULL 的业务必填列（多为父记录外键），
与「严谨的业务逻辑关系」相悖。本脚本将这类列收紧为 NOT NULL，使库约束与 ORM 一致。

安全前提：先校验该列无 NULL 值（有 NULL 则跳过并上报），无 NULL 才执行 ALTER，
故 SET NOT NULL 不会因存量数据失败。相关外键均为 NO ACTION / CASCADE（无 SET NULL），
与 NOT NULL 兼容。

幂等：已 NOT NULL 的列跳过；可重复执行。
用法（backend/ 目录）：
    python scripts/align_columns.py --check   # 只读预览
    python scripts/align_columns.py            # 执行
"""
import argparse
import os
import sys

BACKEND_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if BACKEND_DIR not in sys.path:
    sys.path.insert(0, BACKEND_DIR)

from sqlalchemy import create_engine, text

from app.config.setting import settings

# ORM 声明必填（NOT NULL）、线上库却可空的业务列 → 收紧
COLUMNS = [
    ("bre_pedigree", "combination_id"),
    ("bre_pedigree", "child_code"),
    ("bre_selection_result", "tree_id"),
    ("bre_selection_rule", "target_id"),
    ("bre_prediction_value", "prediction_id"),
    ("bre_treatment", "trial_study_id"),
    ("bre_tree_evaluation", "tree_id"),
    ("bre_trial_study", "trial_id"),
    ("bre_trial_study_entry", "trial_study_id"),
    ("bre_trial_study_entry", "entry_number"),
    ("bre_trial_study_entry", "germplasm_id"),
]


def main() -> None:
    parser = argparse.ArgumentParser(description="业务必填列 NOT NULL 对齐")
    parser.add_argument("--check", action="store_true", help="仅打印将执行的 DDL，不落库")
    args = parser.parse_args()

    engine = create_engine(settings.DB_URI)
    print(f"# 连接: {settings.DB_URI}")

    # 每列用单事务：读锁（count）与 ALTER（ACCESS EXCLUSIVE）在同一事务内升级，
    # 避免跨连接自锁（读连接持锁阻塞另一连接的 ALTER）。
    for t, c in COLUMNS:
        with engine.begin() as tx:
            try:
                row = tx.execute(text(
                    "SELECT is_nullable FROM information_schema.columns "
                    "WHERE table_name=:t AND column_name=:c"
                ), {"t": t, "c": c}).fetchone()
            except Exception as e:  # noqa: BLE001
                print(f"# [跳过] {t}.{c}: 查列失败 {e!s}")
                continue
            if row is None:
                print(f"# [跳过] {t}.{c}: 列不存在")
                continue
            if row[0] == "NO":
                print(f"# [已存在] {t}.{c} 已 NOT NULL")
                continue
            nulls = tx.execute(text(f"SELECT count(*) FROM {t} WHERE {c} IS NULL")).scalar()
            if nulls:
                print(f"# [跳过] {t}.{c}: 存在 {nulls} 个 NULL 值，拒绝收紧")
                continue
            ddl = f"ALTER TABLE {t} ALTER COLUMN {c} SET NOT NULL"
            print("  " + ddl)
            if not args.check:
                try:
                    tx.execute(text(ddl))
                    print("OK  " + ddl)
                except Exception as e:  # noqa: BLE001
                    print("ERR " + ddl + " -> " + str(e)[:160])

    print("# --check 模式：未执行任何变更。" if args.check else "# 完成。")


if __name__ == "__main__":
    main()