import streamlit as st
from moviepy.editor import VideoFileClip, AudioFileClip
from transformers import MarianMTModel, MarianTokenizer
import torch
from gtts import gTTS
import os

# Interface do Streamlit
st.title("Tradutor de Vídeo com Áudio")
uploaded_file = st.file_uploader("Carregue um vídeo para tradução", type=["mp4", "mov", "avi"])

# Função para traduzir o áudio do vídeo
def translate_audio(input_audio_path, output_audio_path, model_name='Helsinki-NLP/opus-mt-en-pt'):
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)

    audio_clip = AudioFileClip(input_audio_path)
    text = "Hello, how are you?"  # Simulação de transcrição
    translated_text = translate_text(text, tokenizer, model)

    tts = gTTS(translated_text, lang='pt')
    tts.save(output_audio_path)

# Função para traduzir o texto
def translate_text(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    with torch.no_grad():
        translated = model.generate(**inputs)
    return tokenizer.decode(translated[0], skip_special_tokens=True)

# Função para substituir o áudio no vídeo
def replace_audio_in_video(input_video_path, input_audio_path, output_video_path):
    video = VideoFileClip(input_video_path)
    new_audio = AudioFileClip(input_audio_path)
    video = video.set_audio(new_audio)
    video.write_videofile(output_video_path, codec="libx264")

# Processar o vídeo quando carregado
if uploaded_file is not None:
    input_video_path = "uploaded_video.mp4"
    output_video_path = "translated_video.mp4"
    
    with open(input_video_path, "wb") as f:
        f.write(uploaded_file.read())

    process_video(input_video_path, output_video_path)
    st.video(output_video_path)
    st.download_button("Baixar Vídeo Traduzido", open(output_video_path, "rb"), "translated_video.mp4")