import os
import subprocess
import sys

# ==========================================
# 1. PREPARAÇÃO DO AMBIENTE
# ==========================================
os.environ["COQUI_TOS_AGREED"] = "1"
os.environ["TORCH_LOAD_WEIGHTS_ONLY"] = "0"

def install_packages():
    print("Verificando/Instalando pacotes necessários...")
    pacotes = [
        "gradio>=4.0.0",
        "librosa",
        "soundfile",
        "TTS==0.22.0",
        "transformers==4.36.2",
        "spaces"
    ]
    subprocess.check_call([sys.executable, "-m", "pip", "install"] + pacotes)
    print("Dependências instaladas com sucesso!")

try:
    import librosa
    import TTS
    import spaces
except ImportError:
    install_packages()

# ==========================================
# 2. IMPORTS PRINCIPAIS
# ==========================================
import gradio as gr
import librosa
import numpy as np
import torch
import spaces
from TTS.api import TTS

# ==========================================
# 3. PATCH DO PYTORCH
# ==========================================
_original_load = torch.load
def _patched_load(*args, **kwargs):
    kwargs["weights_only"] = False
    return _original_load(*args, **kwargs)
torch.load = _patched_load

tts_model = None

# ==========================================
# 4. LÓGICA DO PROCESSAMENTO (COM ZEROGPU)
# ==========================================
# O Hugging Face EXIGE este decorador para liberar a placa de vídeo
@spaces.GPU
def processar_cyr(caminho_referencia, caminho_alvo):
    global tts_model
    
    if caminho_referencia is None or caminho_alvo is None:
        return "Erro: Envie ambos os arquivos de áudio antes de prosseguir."
    
    try:
        # Carrega o modelo apenas quando a GPU for alocada para esta função
        if tts_model is None:
            print("Iniciando motor TTS na GPU...")
            # Força o uso da GPU já que o decorador garantiu o acesso
            tts_model = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=True)
            print("Motor carregado!")

        audio_ref, sr_ref = librosa.load(caminho_referencia, sr=None)
        duracao_ref = librosa.get_duration(y=audio_ref, sr=sr_ref)
        
        audio_alvo, sr_alvo = librosa.load(caminho_alvo, sr=None)
        duracao_alvo = librosa.get_duration(y=audio_alvo, sr=sr_alvo)
        
        resultado = (f"Áudio de Referência processado: Duração de {duracao_ref:.2f} segundos (Taxa: {sr_ref}Hz)\n"
                     f"Áudio Alvo processado: Duração de {duracao_alvo:.2f} segundos (Taxa: {sr_alvo}Hz)\n"
                     "Análise básica concluída com sucesso! (Motor XTTS ativo na memória da GPU)")
                     
    except Exception as e:
        resultado = f"Ocorreu um erro ao processar os áudios: {str(e)}"
        
    return resultado

# ==========================================
# 5. INTERFACE VISUAL
# ==========================================
with gr.Blocks(title="CloneMe") as demo:
    gr.Markdown("# CloneMe")
    
    with gr.Row():
        ref_input = gr.Audio(label="Upload da Referência", type="filepath")
        alvo_input = gr.Audio(label="Upload do Alvo", type="filepath")
    
    btn_next = gr.Button("Next Step")
    output = gr.Textbox(label="Status do Processamento", lines=4)
    
    btn_next.click(
        fn=processar_cyr, 
        inputs=[ref_input, alvo_input], 
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)