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

def create_admin():
    """Cria o usuário admin se não existir"""
    db = get_db_connection()
    try:
        # Verificar se o admin já existe
        admin = db.query(Usuario).filter_by(username='admin').first()
        
        if admin:
            print("Usuário admin já existe")
            
            # Verificar se o arquivo admin_qr.png existe
            if os.path.exists('admin_qr.png'):
                print("Usando OTP existente do arquivo admin_qr.png")
                # Extrair o OTP secret do QR code existente
                with open('admin_qr.png', 'rb') as f:
                    qr_data = f.read()
                    # Aqui você precisaria implementar a lógica para extrair o OTP secret do QR code
                    # Por enquanto, vamos apenas manter o OTP existente
                    return
            else:
                print("Gerando novo OTP para o admin...")
                # Gerar novo OTP
                otp_secret = pyotp.random_base32()
                admin.otp_secret = otp_secret
                db.commit()
        else:
            print("Criando usuário admin...")
            # Criar usuário admin
            admin = Usuario(
                username='admin',
                password='admin123',
                is_admin=True
            )
            admin.email = 'admin@controles.com'
            db.add(admin)
            db.commit()
            
            # Gerar OTP
            otp_secret = pyotp.random_base32()
            admin.otp_secret = otp_secret
            db.commit()
            
            # Atribuir role de Secretário Geral
            admin_role = db.query(Role).filter_by(nivel=Role.SECRETARIO_GERAL).first()
            if admin_role:
                admin.roles.append(admin_role)
                db.commit()
        
        # Gerar QR code
        totp = pyotp.TOTP(otp_secret)
        provisioning_uri = totp.provisioning_uri(admin.username, issuer_name="Sistema de Controles")
        
        qr = qrcode.QRCode(version=1, box_size=10, border=5)
        qr.add_data(provisioning_uri)
        qr.make(fit=True)
        
        img = qr.make_image(fill_color="black", back_color="white")
        
        # Salvar QR code como base64
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        qr_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        # Salvar QR code como arquivo
        img.save('admin_qr.png')
        
        print("\nConfiguração do OTP para o admin:")
        print(f"OTP Secret: {otp_secret}")
        print("\nInstruções:")
        print("1. Use um aplicativo autenticador (como Google Authenticator ou Authy)")
        print("2. Escaneie o QR code ou insira o OTP Secret manualmente")
        print("3. Use o código gerado para fazer login")
        print("\nQR code salvo em 'admin_qr.png'")
        
    except Exception as e:
        print(f"Erro ao criar admin: {str(e)}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == '__main__':
    create_admin()