#!/usr/bin/env python3
"""
更新数据库结构，删除api_keys表，添加roles表，修改users表
"""

from sqlalchemy.orm import Session
from sqlalchemy import text
from app.models.database import engine, Base, SessionLocal
from app.models.models import User, Role
from app.services.user import UserService


def update_db():
    """更新数据库结构"""
    db = SessionLocal()
    
    try:
        # 1. 删除api_keys表
        print("删除api_keys表...")
        try:
            db.execute(text("DROP TABLE IF EXISTS api_keys CASCADE"))
            db.commit()
            print("✅ api_keys表删除成功")
        except Exception as e:
            print(f"⚠️ 删除api_keys表时出错: {e}")
            db.rollback()
        
        # 2. 创建roles表
        print("\n创建roles表...")
        try:
            # 直接执行SQL创建表，避免依赖模型的顺序
            db.execute(text("""
                CREATE TABLE IF NOT EXISTS roles (
                    id VARCHAR PRIMARY KEY,
                    name VARCHAR UNIQUE NOT NULL,
                    description TEXT DEFAULT '',
                    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
                    updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """))
            db.commit()
            print("✅ roles表创建成功")
        except Exception as e:
            print(f"⚠️ 创建roles表时出错: {e}")
            db.rollback()
        
        # 3. 修改users表，添加role_id字段，删除role字段
        print("\n修改users表...")
        try:
            # 检查是否存在role_id字段
            result = db.execute(text("SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'role_id'")).fetchone()
            
            if not result:
                # 添加role_id字段
                db.execute(text("ALTER TABLE users ADD COLUMN role_id VARCHAR"))
                print("✅ 添加role_id字段成功")
            
            # 初始化默认角色
            UserService.init_default_roles(db)
            print("✅ 默认角色初始化成功")
            
            # 获取默认角色
            admin_role = UserService.get_role_by_name(db, "admin")
            user_role = UserService.get_role_by_name(db, "user")
            
            if admin_role and user_role:
                # 更新现有用户的role_id字段
                db.execute(text(f"UPDATE users SET role_id = CASE WHEN role = 'admin' THEN '{admin_role.id}' ELSE '{user_role.id}' END"))
                print("✅ 更新用户role_id字段成功")
            
            # 删除role字段
            result = db.execute(text("SELECT column_name FROM information_schema.columns WHERE table_name = 'users' AND column_name = 'role'")).fetchone()
            if result:
                db.execute(text("ALTER TABLE users DROP COLUMN role"))
                print("✅ 删除role字段成功")
            
            # 添加外键约束
            db.execute(text("ALTER TABLE users ADD CONSTRAINT fk_users_role FOREIGN KEY (role_id) REFERENCES roles(id)"))
            print("✅ 添加外键约束成功")
            
            db.commit()
        except Exception as e:
            print(f"⚠️ 修改users表时出错: {e}")
            db.rollback()
        
        # 4. 检查并创建默认管理员账号
        print("\n检查默认管理员账号...")
        try:
            # 检查是否已存在管理员账号
            admin_user = db.query(User).filter(User.username == "admin").first()
            
            if not admin_user:
                # 获取默认管理员角色
                admin_role = UserService.get_role_by_name(db, "admin")
                if admin_role:
                    # 创建默认管理员账号
                    admin_user = User(
                        id="user-admin",
                        username="admin",
                        email="admin@example.com",
                        password_hash=UserService.get_password_hash("admin123"),
                        role_id=admin_role.id,
                        status="active"
                    )
                    db.add(admin_user)
                    db.commit()
                    db.refresh(admin_user)
                    print("✅ 默认管理员账号创建成功")
                    print(f"用户名: admin")
                    print(f"密码: admin123")
                else:
                    print("❌ 无法创建管理员账号，因为admin角色不存在")
            else:
                print("⚠️ 管理员账号已存在")
                
        except Exception as e:
            print(f"⚠️ 检查管理员账号时出错: {e}")
            db.rollback()
        
        print("\n✅ 数据库结构更新完成")
        
    finally:
        db.close()


if __name__ == "__main__":
    update_db()