import gradio as gr
import whisper
import torch
import os
import ffmpeg
from flask import send_file
import time
from deep_translator import GoogleTranslator  

device = "cuda" if torch.cuda.is_available() else "cpu"
print("Usando dispositivo:", device)

model = whisper.load_model("medium").to(device)  

def extract_audio(video_path, output_audio="temp_audio.wav"):
    ffmpeg.input(video_path).output(output_audio, format="wav", acodec="pcm_s16le", ac=1, ar="16k").run(overwrite_output=True, quiet=True)
    return output_audio

def translate_text(text, target_lang="es"):
    return GoogleTranslator(source="auto", target=target_lang).translate(text)

def transcribe_audio(audio_path, target_language="es"):
    start_time = time.time()
    result = model.transcribe(audio_path, fp16=(device == "cuda"), task="translate", beam_size=5)
    translated_text = translate_text(result["text"], target_language)  
    end_time = time.time()
    print(f"Tiempo de transcripción y traducción: {end_time - start_time:.2f} segundos")
    return translated_text

def save_transcription(text, format, output_name="translation"):
    file_path = f"{output_name}.{format.lower()}"
    with open(file_path, "w", encoding="utf-8") as f:
        f.write(text)
    return file_path  

def process_media(file, format, target_language):
    if not os.path.exists(file):
        return "Error: No se encontró el archivo."
    extension = file.split(".")[-1].lower()
    audio_extensions = ["mp3", "wav", "flac", "ogg", "m4a"]
    video_extensions = ["mp4", "avi", "mov", "mkv"]

    if extension in video_extensions:
        audio_path = extract_audio(file)
    elif extension in audio_extensions:
        audio_path = file
    else:
        return "Error: Formato de archivo no compatible."

    transcription = transcribe_audio(audio_path, target_language)
    file_path = save_transcription(transcription, format)
    
    return transcription, file_path

gui = gr.Interface(
    fn=process_media,
    inputs=[
        gr.File(label="Sube un archivo de audio o vídeo"),
        gr.Radio(["TXT", "DOCX", "SRT", "JSON"], label="Formato de salida"),
        gr.Dropdown(["es", "en", "fr", "de", "it"], label="Idioma de traducción", value="es")
    ],
    outputs=["text", "file"],  
    title="Transcriptor y Traductor de Audio y Vídeo",
    description="Sube un archivo, transcribe y traduce a cualquier idioma con alta precisión."
)

gui.launch()