import os
import subprocess
import soundfile as sf
import librosa
from transformers import WhisperProcessor, WhisperForConditionalGeneration, MarianMTModel, MarianTokenizer
from gtts import gTTS
import gradio as gr

# Configurações
WHISPER_MODEL = "openai/whisper-medium"
TRANSLATION_MODEL = "Helsinki-NLP/opus-mt-en-pt"
TEMP_AUDIO_FILE = "temp_audio.wav"
TEMP_TRANSLATED_AUDIO = "translated_audio.mp3"
OUTPUT_VIDEO = "output_video.mp4"

def extract_audio(video_path: str, output_audio_path: str) -> None:
    """
    Extrai o áudio de um vídeo usando FFmpeg.
    """
    try:
        command = [
            "ffmpeg", "-i", video_path, "-q:a", "0", "-map", "a", output_audio_path, "-y"
        ]
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        raise Exception(f"Erro ao extrair áudio: {e}")

def transcribe_audio(audio_path: str) -> str:
    """
    Transcreve o áudio para texto usando o modelo Whisper.
    """
    try:
        # Carregar o áudio
        audio, sample_rate = sf.read(audio_path)

        # Verificar se o áudio está no formato correto
        if len(audio.shape) > 1:  # Se for stereo, converter para mono
            audio = audio.mean(axis=1)

        # Redimensionar para 16 kHz, se necessário
        if sample_rate != 16000:
            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)

        # Processar o áudio
        processor = WhisperProcessor.from_pretrained(WHISPER_MODEL)
        model = WhisperForConditionalGeneration.from_pretrained(WHISPER_MODEL)
        input_features = processor(audio, sampling_rate=16000, return_tensors="pt").input_features

        # Gerar transcrição (forçando o idioma para inglês)
        predicted_ids = model.generate(input_features, language="en", task="transcribe")
        transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
        return transcription
    except Exception as e:
        raise Exception(f"Erro ao transcrever áudio: {e}")

def translate_text(text: str, source_lang: str = "en", target_lang: str = "pt") -> str:
    """
    Traduz o texto para o idioma desejado usando o modelo MarianMT.
    """
    try:
        model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
        tokenizer = MarianTokenizer.from_pretrained(model_name)
        model = MarianMTModel.from_pretrained(model_name)
        translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
        translation = tokenizer.decode(translated[0], skip_special_tokens=True)
        return translation
    except Exception as e:
        raise Exception(f"Erro ao traduzir texto: {e}")

def text_to_speech(text: str, output_audio_path: str, lang: str = "pt") -> None:
    """
    Converte texto em áudio usando gTTS.
    """
    try:
        tts = gTTS(text, lang=lang)
        tts.save(output_audio_path)
    except Exception as e:
        raise Exception(f"Erro ao gerar áudio: {e}")

def merge_audio_video(video_path: str, audio_path: str, output_path: str) -> None:
    """
    Combina o áudio traduzido com o vídeo original usando FFmpeg.
    """
    try:
        command = [
            "ffmpeg", "-i", video_path, "-i", audio_path, "-c:v", "copy", "-map", "0:v:0", "-map", "1:a:0", "-shortest", output_path, "-y"
        ]
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        raise Exception(f"Erro ao combinar áudio e vídeo: {e}")

def translate_video(video_path: str) -> str:
    """
    Função principal que orquestra a tradução do vídeo.
    """
    try:
        # 1. Extrair áudio do vídeo
        extract_audio(video_path, TEMP_AUDIO_FILE)

        # 2. Transcrever áudio para texto
        transcription = transcribe_audio(TEMP_AUDIO_FILE)
        print(f"Transcrição: {transcription}")

        # 3. Traduzir texto para português
        translation = translate_text(transcription)
        print(f"Tradução: {translation}")

        # 4. Converter texto traduzido em áudio
        text_to_speech(translation, TEMP_TRANSLATED_AUDIO)

        # 5. Combinar áudio traduzido com vídeo original
        merge_audio_video(video_path, TEMP_TRANSLATED_AUDIO, OUTPUT_VIDEO)

        # 6. Limpar arquivos temporários
        os.remove(TEMP_AUDIO_FILE)
        os.remove(TEMP_TRANSLATED_AUDIO)

        return OUTPUT_VIDEO
    except Exception as e:
        print(f"Erro durante o processamento: {e}")
        raise

# Interface Gradio
def gradio_interface(video):
    try:
        output_video = translate_video(video)
        return output_video
    except Exception as e:
        return f"Erro: {e}"

# Configuração da interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Video(label="Upload do Vídeo"),
    outputs=gr.Video(label="Vídeo Traduzido"),
    title="Tradutor de Vídeos para Português",
    description="Faça upload de um vídeo em qualquer idioma e receba o vídeo com áudio traduzido para português."
)

# Iniciar a interface
iface.launch()