import pymysql
import sys

config = {
    'host': 'localhost',
    'port': 3306,
    'user': 'root',
    'password': 'rootx77',
    'database': 'ai',
    'charset': 'utf8mb4'
}

print("开始导入数据库...")

try:
    conn = pymysql.connect(**config)
    cursor = conn.cursor()
    
    # 读取 SQL 文件
    with open('数据库.sql', 'r', encoding='utf8') as f:
        sql_content = f.read()
    
    # 分割 SQL 语句
    statements = sql_content.split(';')
    
    total = len(statements)
    success = 0
    
    for i, statement in enumerate(statements):
        statement = statement.strip()
        if statement and not statement.startswith('--'):
            try:
                cursor.execute(statement)
                success += 1
                if i % 10 == 0:
                    print(f"进度: {i}/{total}")
            except Exception as e:
                # 忽略某些错误（如表已存在）
                if 'already exists' not in str(e):
                    pass
    
    conn.commit()
    
    # 检查导入结果
    cursor.execute("SHOW TABLES")
    tables = cursor.fetchall()
    
    print(f"\n✅ 导入完成！")
    print(f"成功执行: {success}/{total} 条语句")
    print(f"数据库中共有 {len(tables)} 张表")
    
    # 检查关键表
    key_tables = ['nf_user', 'nf_lovers', 'nf_chat_message']
    print("\n检查关键表:")
    for table_name in key_tables:
        cursor.execute(f"SHOW TABLES LIKE '{table_name}'")
        if cursor.fetchone():
            print(f"  ✅ {table_name}")
        else:
            print(f"  ❌ {table_name}")
    
    cursor.close()
    conn.close()
    
except Exception as e:
    print(f"❌ 错误: {e}")
    sys.exit(1)