"""
迁移脚本：为 family_member_info 表添加 data_source 字段
- 新增 data_source VARCHAR(20) NOT NULL DEFAULT 'backend'
- 将所有现有记录赋值为 'backend'
- 取值约定：'backend' | 'miniprogram' | 'pdf_ai'
"""

import pymysql

DB_CONFIG = {
    "host": "rm-f8ze60yirdj8786u2wo.mysql.rds.aliyuncs.com",
    "port": 3306,
    "user": "root",
    "password": "csqz@20255",
    "db": "csqz-client",
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor
}

def run():
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cursor:
            # 1. 检查字段是否已存在
            cursor.execute("SHOW COLUMNS FROM family_member_info LIKE 'data_source'")
            exists = cursor.fetchone()

            if exists:
                print("[Migration] data_source 字段已存在，跳过 ALTER TABLE")
            else:
                print("[Migration] 添加 data_source 字段...")
                cursor.execute("""
                    ALTER TABLE family_member_info
                    ADD COLUMN data_source VARCHAR(20) NOT NULL DEFAULT 'backend'
                    COMMENT '录入来源: backend=后台手工, miniprogram=小程序, pdf_ai=PDF智能识别'
                """)
                conn.commit()
                print("[Migration] 字段添加成功")

            # 2. 将所有现有记录显式赋值为 'backend'（DEFAULT 已覆盖，这里做二次保障）
            cursor.execute("SELECT COUNT(*) as total FROM family_member_info WHERE data_source != 'backend' OR data_source IS NULL")
            row = cursor.fetchone()
            not_backend = row['total']

            if not_backend > 0:
                print(f"[Migration] 发现 {not_backend} 条记录需要更新为 'backend'...")
                cursor.execute("UPDATE family_member_info SET data_source = 'backend' WHERE data_source != 'backend' OR data_source IS NULL")
                conn.commit()
                print(f"[Migration] 已将 {cursor.rowcount} 条记录更新为 'backend'")
            else:
                print("[Migration] 所有现有记录已为 'backend'，无需更新")

            # 3. 验证结果
            cursor.execute("SELECT data_source, COUNT(*) as cnt FROM family_member_info GROUP BY data_source")
            stats = cursor.fetchall()
            print("\n[Migration] 当前 data_source 分布：")
            for row in stats:
                print(f"  {row['data_source']}: {row['cnt']} 条")

        print("\n[Migration] 迁移完成 ✓")
    except Exception as e:
        conn.rollback()
        print(f"[Migration] 错误: {e}")
        raise
    finally:
        conn.close()

if __name__ == '__main__':
    run()