import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from subprocess import run
from faster_whisper import WhisperModel
import json
import tempfile

# Carregar mapeamento de idiomas
with open('language_codes.json', 'r') as f:
    lang_codes = json.load(f)

# Inicializar modelos
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
whisper_model = WhisperModel("large-v2", device="cuda", compute_type="float16")

def process_video(radio, video, target_language, use_wav2lip):
    # 1. Extrair áudio
    audio_file = tempfile.NamedTemporaryFile(suffix=".wav").name
    run(["ffmpeg", "-i", video.name, audio_file])

    # 2. Transcrição
    segments, _ = whisper_model.transcribe(audio_file)
    transcript = " ".join([segment.text for segment in segments])

    # 3. Tradução
    flores_code = lang_codes.get(target_language, "eng_Latn")
    inputs = tokenizer(transcript, return_tensors="pt")
    translated_tokens = model.generate(**inputs, forced_bos_token_id=tokenizer.lang_code_to_id[flores_code], max_length=100)
    translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]

    # 4. Criar arquivo de legenda
    subtitle_file = tempfile.NamedTemporaryFile(suffix=".srt", delete=False).name
    with open(subtitle_file, "w") as f:
        f.write("1\n00:00:00,000 --> 00:00:10,000\n" + translated_text)

    # 5. Incorporar legenda
    output_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    run(["ffmpeg", "-i", video.name, "-vf", f"subtitles={subtitle_file}", output_video])

    return output_video

# Interface Gradio
iface = gr.Interface(
    fn=process_video,
    inputs=[
        gr.Radio(["Upload", "Record"], value="Upload", show_label=False),
        gr.Video(),
        gr.Dropdown(choices=list(lang_codes.keys()), label="Target Language for Dubbing", value="English"),
        gr.Checkbox(label="Video has a close-up face. Use Wav2lip.", value=False)
    ],
    outputs=gr.Video(),
    live=False,
    title="AI Video Dubbing"
)

iface.launch()