"""
数据库迁移脚本：为 users 表添加 email 和 qq_openid 字段
运行: python migrate_add_email_qq_fields.py
"""
import asyncio
import aiomysql
from config import settings
import sys


async def migrate_add_email_qq_fields():
    """为 users 表添加 email 和 qq_openid 字段"""
    
    try:
        # 创建数据库连接
        conn = await aiomysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD,
            db=settings.MYSQL_DATABASE,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            # 检查 users 表是否存在
            await cursor.execute("""
                SELECT COUNT(*) FROM information_schema.tables 
                WHERE table_schema = %s AND table_name = 'users'
            """, (settings.MYSQL_DATABASE,))
            table_exists = (await cursor.fetchone())[0] > 0
            
            if not table_exists:
                print("❌ users 表不存在，请先创建表")
                conn.close()
                return False
            
            # 检查是否已存在 email 列
            await cursor.execute("""
                SELECT COUNT(*) FROM information_schema.columns 
                WHERE table_schema = %s 
                AND table_name = 'users' 
                AND column_name = 'email'
            """, (settings.MYSQL_DATABASE,))
            has_email = (await cursor.fetchone())[0] > 0
            
            # 检查是否已存在 qq_openid 列
            await cursor.execute("""
                SELECT COUNT(*) FROM information_schema.columns 
                WHERE table_schema = %s 
                AND table_name = 'users' 
                AND column_name = 'qq_openid'
            """, (settings.MYSQL_DATABASE,))
            has_qq_openid = (await cursor.fetchone())[0] > 0
            
            if has_email and has_qq_openid:
                print("✅ email 和 qq_openid 字段已存在，无需迁移")
                conn.close()
                return True
            
            # 添加 email 字段（如果不存在）
            if not has_email:
                print("🔄 添加 email 字段...")
                await cursor.execute("""
                    ALTER TABLE users 
                    ADD COLUMN email VARCHAR(200) UNIQUE NULL COMMENT '邮箱' AFTER phone
                """)
                await cursor.execute("""
                    CREATE INDEX idx_email ON users(email)
                """)
                print("✅ email 字段添加成功")
            
            # 添加 qq_openid 字段（如果不存在）
            if not has_qq_openid:
                print("🔄 添加 qq_openid 字段...")
                await cursor.execute("""
                    ALTER TABLE users 
                    ADD COLUMN qq_openid VARCHAR(100) UNIQUE NULL COMMENT 'QQ OpenID' AFTER wechat_openid
                """)
                await cursor.execute("""
                    CREATE INDEX idx_qq_openid ON users(qq_openid)
                """)
                print("✅ qq_openid 字段添加成功")
            
            await conn.commit()
            print("✅ 迁移完成")
        
        conn.close()
        return True
        
    except Exception as e:
        import traceback
        print(f"❌ 迁移失败: {e}")
        traceback.print_exc()
        return False


if __name__ == "__main__":
    # 确认是否继续
    if len(sys.argv) > 1 and sys.argv[1] == "--force":
        confirm = "y"
    else:
        confirm = input("\n是否继续迁移 users 表添加 email 和 qq_openid 字段? (y/n): ")
    
    if confirm.lower() != 'y':
        sys.exit(0)
    
    # 执行迁移
    success = asyncio.run(migrate_add_email_qq_fields())
    
    if not success:
        sys.exit(1)