import pymysql

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

def update_db():
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cursor:
            # 1. Update genealogy_records
            print("Updating genealogy_records table...")
            try:
                cursor.execute("ALTER TABLE genealogy_records ADD COLUMN ai_status INT DEFAULT 0 COMMENT '0:未处理, 1:处理中, 2:成功, 3:失败'")
                print("Added ai_status column.")
            except Exception as e:
                print(f"Skipping ai_status: {e}")

            try:
                cursor.execute("ALTER TABLE genealogy_records ADD COLUMN ai_result LONGTEXT COMMENT 'AI解析结果JSON'")
                print("Added ai_result column.")
            except Exception as e:
                print(f"Skipping ai_result: {e}")

            # 2. Update family_member_info
            print("Updating family_member_info table...")
            try:
                cursor.execute("ALTER TABLE family_member_info ADD COLUMN source_record_id INT COMMENT '关联的家谱记录ID'")
                print("Added source_record_id column.")
            except Exception as e:
                print(f"Skipping source_record_id: {e}")

        conn.commit()
        print("Database update complete.")
    except Exception as e:
        print(f"Error updating database: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    update_db()