#!/usr/bin/env python
"""
添加用户表的学校和专业字段
"""

import asyncio
import sys
from pathlib import Path
import aiomysql
import os

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

from app.core.config import settings

# SQL脚本
ADD_FIELDS_SQL = """
-- 添加 school 和 major 字段到 users 表
ALTER TABLE users 
ADD COLUMN school VARCHAR(100) COMMENT '学校' AFTER bio,
ADD COLUMN major VARCHAR(100) COMMENT '专业' AFTER school;
"""

async def execute_sql():
    """执行SQL脚本"""
    try:
        # 从环境变量或配置中获取数据库连接信息
        # 优先使用环境变量
        database_url = os.getenv("DATABASE_URL", settings.DATABASE_URL)
        
        # 解析数据库连接字符串
        # 格式: mysql+aiomysql://root:root@localhost:3306/kaopeilian
        if database_url.startswith("mysql+aiomysql://"):
            url = database_url.replace("mysql+aiomysql://", "")
        else:
            url = database_url
            
        # 解析连接参数
        auth_host = url.split("@")[1]
        user_pass = url.split("@")[0]
        host_port_db = auth_host.split("/")
        host_port = host_port_db[0].split(":")
        
        user = user_pass.split(":")[0]
        password = user_pass.split(":")[1]
        host = host_port[0]
        port = int(host_port[1]) if len(host_port) > 1 else 3306
        database = host_port_db[1].split("?")[0] if len(host_port_db) > 1 else "kaopeilian"
        
        print(f"连接数据库: {host}:{port}/{database}")
        
        # 创建数据库连接
        conn = await aiomysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            db=database,
            charset='utf8mb4'
        )
        
        async with conn.cursor() as cursor:
            # 检查字段是否已存在
            check_sql = """
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = %s 
              AND TABLE_NAME = 'users' 
              AND COLUMN_NAME IN ('school', 'major')
            """
            await cursor.execute(check_sql, (database,))
            existing_columns = [row[0] for row in await cursor.fetchall()]
            
            if 'school' in existing_columns or 'major' in existing_columns:
                print("字段已存在，跳过添加")
                if 'school' in existing_columns:
                    print("- school 字段已存在")
                if 'major' in existing_columns:
                    print("- major 字段已存在")
            else:
                # 执行SQL脚本
                await cursor.execute(ADD_FIELDS_SQL)
                await conn.commit()
                print("成功添加 school 和 major 字段到 users 表")
            
            # 显示表结构
            show_sql = "DESC users"
            await cursor.execute(show_sql)
            columns = await cursor.fetchall()
            
            print("\n当前 users 表结构：")
            print("-" * 80)
            for col in columns:
                print(f"{col[0]:20} {col[1]:30} {'NULL' if col[2] == 'YES' else 'NOT NULL':10} {col[5] or ''}")
        
        conn.close()
        
    except Exception as e:
        print(f"执行失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    print("开始添加 school 和 major 字段...")
    asyncio.run(execute_sql())
    print("\n执行完成！")