"""数据治理 migration: 入库必检约束 + 元数据字段补充"""
import pymysql
import os
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("data-governance")

DB_USER = os.getenv("CMA_DB_USER", "cma_user")
DB_PASS = os.getenv("CMA_DB_PASS", "cma_pass_2026")
DB_HOST = os.getenv("CMA_DB_HOST", "127.0.0.1")
DB_PORT = int(os.getenv("CMA_DB_PORT", "3306"))
DB_NAME = os.getenv("CMA_DB_NAME", "cma")

conn = pymysql.connect(
    host=DB_HOST, port=DB_PORT, user=DB_USER, password=DB_PASS, database=DB_NAME,
    charset="utf8mb4", cursorclass=pymysql.cursors.DictCursor,
)
cur = conn.cursor()


def run():
    logger.info("=== 数据治理 Migration 开始 ===")

    # ── 0. 添加缺失列 ──
    logger.info("[步骤0] 检查并补充缺失列...")

    for col_name, col_def in [
        ("data_source", "ALTER TABLE kpi_definitions ADD COLUMN data_source VARCHAR(500) DEFAULT NULL COMMENT '数据来源' AFTER data_source_type"),
        ("data_owner", "ALTER TABLE kpi_definitions ADD COLUMN data_owner VARCHAR(100) DEFAULT NULL COMMENT '数据责任人' AFTER data_source"),
    ]:
        cur.execute("SHOW COLUMNS FROM kpi_definitions LIKE %s", (col_name,))
        if not cur.fetchone():
            logger.info(f"  添加 {col_name} 列...")
            cur.execute(col_def)
            conn.commit()
            logger.info(f"  {col_name} 列已添加")
        else:
            logger.info(f"  {col_name} 列已存在")

    # ── 1. 修复空值 ──
    logger.info("[步骤1] 修复空值...")

    for col, default, label in [
        ("target_value", "0", "NULL target_value"),
        ("unit", "'-'", "NULL/empty unit"),
        ("formula", "'待补充'", "NULL/empty formula"),
        ("data_source", "'待补充'", "NULL/empty data_source"),
        ("data_owner", "'待指定'", "NULL/empty data_owner"),
    ]:
        if col in ("target_value",):
            r = cur.execute(f"SELECT COUNT(*) as cnt FROM kpi_definitions WHERE {col} IS NULL")
        else:
            r = cur.execute(f"SELECT COUNT(*) as cnt FROM kpi_definitions WHERE {col} IS NULL OR {col} = ''")
        row = cur.fetchone()
        cnt = row["cnt"] if row else 0
        logger.info(f"  {label}: {cnt} 条")
        if cnt > 0:
            if col in ("target_value",):
                cur.execute(f"UPDATE kpi_definitions SET {col} = {default} WHERE {col} IS NULL")
            else:
                cur.execute(f"UPDATE kpi_definitions SET {col} = {default} WHERE {col} IS NULL OR {col} = ''")
            logger.info(f"  已修复 {cur.rowcount} 条")

    conn.commit()

    # ── 2. 修改列约束为 NOT NULL ──
    logger.info("[步骤2] 修改列约束...")

    alters = [
        ("target_value", "DECIMAL(15,2) NOT NULL DEFAULT 0"),
        ("unit", "VARCHAR(50) NOT NULL DEFAULT '-'"),
        ("formula", "TEXT NOT NULL"),
        ("data_source", "VARCHAR(500) NOT NULL DEFAULT '待补充'"),
        ("data_owner", "VARCHAR(100) NOT NULL DEFAULT '待指定'"),
    ]
    for col, col_type in alters:
        col_comment = {
            "target_value": "目标值", "unit": "单位", "formula": "计算公式",
            "data_source": "数据来源", "data_owner": "数据责任人",
        }[col]
        try:
            cur.execute(f"ALTER TABLE kpi_definitions MODIFY {col} {col_type} COMMENT '{col_comment}'")
            logger.info(f"  {col} → {col_type}")
        except Exception as e:
            logger.warning(f"  {col} 修改失败: {e}")

    conn.commit()

    # ── 3. 验证 ──
    logger.info("[步骤3] 验证约束...")
    cur.execute("DESCRIBE kpi_definitions")
    for col in cur.fetchall():
        if col['Field'] in ('target_value', 'unit', 'formula', 'data_source', 'data_owner'):
            logger.info(f"  {col['Field']}: Null={col['Null']}, Default={col['Default']}, Type={col['Type']}")

    logger.info("=== 数据治理 Migration 完成 ===")


if __name__ == "__main__":
    try:
        run()
    finally:
        cur.close()
        conn.close()