import os
import zipfile
from tqdm import tqdm
from faster_whisper import WhisperModel
from datetime import timedelta
import srt
from transformers import pipeline
import torch
import gradio as gr

OUTPUT_DIR = "outputs"
LEGENDAS_DIR = os.path.join(OUTPUT_DIR, "legendas")
os.makedirs(LEGENDAS_DIR, exist_ok=True)

IDIOMAS = {
    "Alemão": "de", "Cebuano": "ceb", "Chinês": "zh", "Coreano": "ko",
    "Espanhol": "es", "Francês": "fr", "Inglês": "en", "Italiano": "it",
    "Japonês": "ja", "Português": "pt", "Russo": "ru", "Tagalog": "tl"
}

print("Carregando modelos de IA...")
try:
    MODELO_WHISPER = "medium"
    device = "cpu"
    compute_type = "int8"
    model_whisper = WhisperModel(MODELO_WHISPER, device=device, compute_type=compute_type)
    model_translator = pipeline("translation", model="facebook/m2m100_418M")
    print("Modelos carregados com sucesso!")
except Exception as e:
    print(f"Erro ao carregar os modelos: {e}")
    model_whisper = None
    model_translator = None

def gerar_legenda_interface(arquivos_path, idioma_origem, idioma_destino, progresso=gr.Progress(track_tqdm=True)):
    if not arquivos_path: raise gr.Error("Por favor, envie pelo menos um arquivo.")
    if not idioma_origem or not idioma_destino: raise gr.Error("Por favor, selecione os idiomas.")

    lang_origem_code = IDIOMAS[idioma_origem]
    lang_destino_code = IDIOMAS[idioma_destino]
    
    if os.path.exists(LEGENDAS_DIR):
        for f in os.listdir(LEGENDAS_DIR): os.remove(os.path.join(LEGENDAS_DIR, f))

    for caminho_video in tqdm(arquivos_path, desc="Processando arquivos"):
        nome_base = os.path.splitext(os.path.basename(caminho_video))[0]
        segmentos, _ = model_whisper.transcribe(caminho_video, language=lang_origem_code, beam_size=5, log_progress=False)
        
        segmentos_lista = list(segmentos)
        textos_originais = [seg.text.strip() for seg in segmentos_lista]

        if lang_origem_code != lang_destino_code:
            textos_finais = model_translator(textos_originais, src_lang=lang_origem_code, tgt_lang=lang_destino_code)
            textos_finais = [t['translation_text'] for t in textos_finais]
        else: textos_finais = textos_originais

        legendas = []
        for i, seg in enumerate(segmentos_lista):
            legendas.append(srt.Subtitle(index=i+1, start=timedelta(seconds=seg.start), end=timedelta(seconds=seg.end), content=textos_finais[i]))

        nome_srt = f"{nome_base}_{lang_destino_code}.srt"
        caminho_srt = os.path.join(LEGENDAS_DIR, nome_srt)
        with open(caminho_srt, "w", encoding="utf-8") as f: f.write(srt.compose(legendas))

    zip_path = os.path.join(OUTPUT_DIR, "legendas_geradas.zip")
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for file in os.listdir(LEGENDAS_DIR): zipf.write(os.path.join(LEGENDAS_DIR, file), arcname=file)
    return zip_path

with gr.Blocks(theme=gr.themes.Default(), title="Gerador de Legendas") as interface:
    gr.Markdown("# 🎬 Gerador Automático de Legendas e Traduções")
    with gr.Row():
        with gr.Column(scale=1):
            input_arquivos = gr.File(label="1. Envie seus arquivos de vídeo ou áudio", file_count="multiple", type="filepath")
            input_idioma_origem = gr.Dropdown(choices=sorted(list(IDIOMAS.keys())), label="2. Idioma original do áudio", value="Inglês")
            input_idioma_destino = gr.Dropdown(choices=sorted(list(IDIOMAS.keys())), label="3. Idioma desejado para a legenda", value="Português")
            botao_gerar = gr.Button("✨ Gerar Legendas!", variant="primary")
        with gr.Column(scale=1):
            output_legenda = gr.File(label="Resultado para Download (Arquivo .ZIP)", interactive=False)
    botao_gerar.click(fn=gerar_legenda_interface, inputs=[input_arquivos, input_idioma_origem, input_idioma_destino], outputs=[output_legenda])

interface.launch()