# -*- coding: utf-8 -*-
"""
在测试库 test_etl_feiqiu 执行员工档案建表迁移脚本。

迁移脚本：db/etl_feiqiu/migrations/2026-02-22__add_staff_info_tables.sql
目标表：ods.staff_info_master, dwd.dim_staff, dwd.dim_staff_ex

使用方式：
    python scripts/ops/run_migration_staff_info.py
"""

import os
import sys
from pathlib import Path

from dotenv import load_dotenv
import psycopg2

# 加载根 .env
_ROOT = Path(__file__).resolve().parents[2]
load_dotenv(_ROOT / ".env", override=False)

DSN = os.getenv("TEST_DB_DSN")
if not DSN:
    print("ERROR: TEST_DB_DSN 未配置，请在根 .env 中设置")
    sys.exit(1)

MIGRATION_FILE = (
    _ROOT / "db" / "etl_feiqiu" / "migrations"
    / "2026-02-22__add_staff_info_tables.sql"
)

# 需要创建的三张表
TABLES = [
    ("ods", "staff_info_master"),
    ("dwd", "dim_staff"),
    ("dwd", "dim_staff_ex"),
]


def tables_exist(conn) -> dict[str, bool]:
    """检查目标表是否已存在，返回 {schema.table: bool}"""
    cur = conn.cursor()
    result = {}
    for schema, table in TABLES:
        cur.execute("""
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = %s AND table_name = %s
        """, (schema, table))
        result[f"{schema}.{table}"] = cur.fetchone() is not None
    cur.close()
    return result


def execute_migration(conn) -> bool:
    """执行迁移脚本，返回是否成功"""
    sql = MIGRATION_FILE.read_text(encoding="utf-8")

    # 去掉注释中的回滚部分
    main_lines = []
    in_rollback = False
    for line in sql.split("\n"):
        stripped = line.strip()
        if stripped.startswith("-- ====") and "回滚" in stripped:
            in_rollback = True
        if not in_rollback:
            main_lines.append(line)

    main_sql = "\n".join(main_lines).strip()
    if not main_sql:
        print("⚠️ 迁移脚本为空，跳过")
        return False

    try:
        cur = conn.cursor()
        cur.execute(main_sql)
        cur.close()
        print("✅ 迁移脚本执行成功")
        return True
    except Exception as e:
        print(f"❌ 迁移脚本执行失败: {e}")
        return False


def verify(conn) -> bool:
    """验证建表结果"""
    cur = conn.cursor()
    checks = []

    # 1. 三张表都存在
    for schema, table in TABLES:
        cur.execute("""
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = %s AND table_name = %s
        """, (schema, table))
        checks.append((f"表 {schema}.{table} 存在", cur.fetchone() is not None))

    # 2. ods.staff_info_master 关键字段
    cur.execute("""
        SELECT column_name FROM information_schema.columns
        WHERE table_schema = 'ods' AND table_name = 'staff_info_master'
    """)
    ods_cols = {r[0] for r in cur.fetchall()}
    ods_required = {"id", "staff_name", "mobile", "content_hash", "payload", "tenant_id", "site_id"}
    missing_ods = ods_required - ods_cols
    checks.append((f"ODS 关键字段完整（{len(ods_cols)} 列）", len(missing_ods) == 0))
    if missing_ods:
        print(f"    ODS 缺失字段: {missing_ods}")

    # 3. dwd.dim_staff 主键包含 staff_id + scd2_start_time
    cur.execute("""
        SELECT column_name FROM information_schema.columns
        WHERE table_schema = 'dwd' AND table_name = 'dim_staff'
    """)
    dwd_cols = {r[0] for r in cur.fetchall()}
    dwd_required = {"staff_id", "staff_name", "scd2_start_time", "scd2_end_time", "scd2_is_current"}
    missing_dwd = dwd_required - dwd_cols
    checks.append((f"DWD 主表关键字段完整（{len(dwd_cols)} 列）", len(missing_dwd) == 0))
    if missing_dwd:
        print(f"    DWD 主表缺失字段: {missing_dwd}")

    # 4. dwd.dim_staff_ex 关键字段
    cur.execute("""
        SELECT column_name FROM information_schema.columns
        WHERE table_schema = 'dwd' AND table_name = 'dim_staff_ex'
    """)
    ex_cols = {r[0] for r in cur.fetchall()}
    ex_required = {"staff_id", "rank_name", "shop_name", "scd2_start_time"}
    missing_ex = ex_required - ex_cols
    checks.append((f"DWD 扩展表关键字段完整（{len(ex_cols)} 列）", len(missing_ex) == 0))
    if missing_ex:
        print(f"    DWD 扩展表缺失字段: {missing_ex}")

    cur.close()

    print("\n" + "=" * 50)
    print("建表验证结果")
    print("=" * 50)
    all_ok = True
    for name, ok in checks:
        status = "✅" if ok else "❌"
        print(f"  {status} {name}")
        if not ok:
            all_ok = False

    return all_ok


def main():
    dsn_display = DSN.split("@")[1] if "@" in DSN else DSN
    print(f"连接测试库: {dsn_display}")
    print(f"迁移脚本: {MIGRATION_FILE.name}\n")

    if not MIGRATION_FILE.exists():
        print(f"ERROR: 迁移脚本不存在: {MIGRATION_FILE}")
        sys.exit(1)

    conn = psycopg2.connect(DSN)
    conn.autocommit = True

    # 检查表是否已存在
    existing = tables_exist(conn)
    all_exist = all(existing.values())

    if all_exist:
        print("ℹ️ 所有目标表已存在，跳过建表")
    else:
        for name, exists in existing.items():
            if exists:
                print(f"  ℹ️ {name} 已存在")
            else:
                print(f"  📋 {name} 待创建")
        if not execute_migration(conn):
            conn.close()
            sys.exit(1)

    # 验证
    all_ok = verify(conn)
    conn.close()

    if all_ok:
        print("\n✅ 员工档案建表迁移完成，所有验证通过")
    else:
        print("\n⚠️ 部分验证未通过，请检查")
        sys.exit(1)


if __name__ == "__main__":
    main()