from functions.database import get_db_connection, Usuario
from functions.rbac import Role
import pyotp
import qrcode
import os
import base64
from io import BytesIO

def create_test_users():
    """Cria usuários de teste se não existirem"""
    db = get_db_connection()
    try:
        # Usuários de teste
        test_users = [
            {
                'username': 'teste',
                'password': 'admin123',  # Mesma senha do admin
                'email': 'teste@controles.com',
                'is_admin': True
            },
            {
                'username': 'aligner',
                'password': 'Test123!@#',
                'email': 'aligner@controles.com',
                'is_admin': False
            },
            {
                'username': 'tester',
                'password': 'Test123!@#',
                'email': 'tester@controles.com',
                'is_admin': False
            },
            {
                'username': 'deployer',
                'password': 'Test123!@#',
                'email': 'deployer@controles.com',
                'is_admin': False
            }
        ]
        
        # Obter o OTP secret do admin se existir
        admin = db.query(Usuario).filter_by(username='admin').first()
        admin_otp_secret = admin.otp_secret if admin else None
        
        for user_data in test_users:
            # Verificar se o usuário já existe
            user = db.query(Usuario).filter_by(username=user_data['username']).first()
            
            if not user:
                print(f"Criando usuário {user_data['username']}...")
                # Criar usuário
                user = Usuario(
                    username=user_data['username'],
                    password=user_data['password'],
                    is_admin=user_data['is_admin']
                )
                user.email = user_data['email']
                db.add(user)
                db.commit()
                
                # Se for o usuário teste, usar o mesmo OTP do admin
                if user_data['username'] == 'teste' and admin_otp_secret:
                    user.otp_secret = admin_otp_secret
                    db.commit()
                else:
                    # Gerar novo OTP para outros usuários
                    otp_secret = pyotp.random_base32()
                    user.otp_secret = otp_secret
                    db.commit()
                
                # Atribuir role de Secretário Geral para o usuário teste
                if user_data['username'] == 'teste':
                    admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
                    if admin_role:
                        user.roles.append(admin_role)
                        db.commit()
                
                print(f"Usuário {user_data['username']} criado com sucesso!")
            else:
                print(f"Usuário {user_data['username']} já existe")
                
                # Se for o usuário teste e não tiver o OTP do admin, atualizar
                if user_data['username'] == 'teste' and admin_otp_secret and user.otp_secret != admin_otp_secret:
                    user.otp_secret = admin_otp_secret
                    db.commit()
                    print(f"OTP do usuário teste atualizado para o mesmo do admin")
                
                # Verificar se o usuário teste tem a role de Secretário Geral
                if user_data['username'] == 'teste':
                    admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
                    if admin_role and admin_role not in user.roles:
                        user.roles.append(admin_role)
                        db.commit()
                        print(f"Role de Secretário Geral atribuída ao usuário teste")
    
    except Exception as e:
        print(f"Erro ao criar usuários de teste: {str(e)}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == '__main__':
    create_test_users() 