import gradio as gr
import whisper
import os

# Modelos disponibles de Whisper
MODEL_OPTIONS = ["tiny", "base", "small", "medium", "large"]

# Cargar el modelo por defecto
model = whisper.load_model("base")

# Idiomas disponibles en el menú desplegable (nombre visible y código real)
IDIOMAS_DISPONIBLES = [
    ("Español", "es"),
    ("Inglés", "en"),
    ("Catalán", "ca"),
    ("Francés", "fr"),
    ("Alemán", "de"),
    ("Italiano", "it"),
    ("Portugués", "pt"),
    ("Gallego", "gl"),
    ("Euskera", "eu"),
    ("Chino", "zh"),
    ("Japonés", "ja"),
]

def transcribir(audio_paths, model_name, idioma_nombre, traducir):
    """
    Función que recibe uno o más archivos de audio, ejecuta la transcripción,
    permite elegir el modelo de Whisper, especificar el idioma y traducir la transcripción.
    Guarda los resultados en archivos .txt y .srt, y devuelve enlaces de descarga.
    """

    global model
    if not model_name:
        gr.Warning("⚠️ Por favor, selecciona un modelo de Whisper antes de continuar.")
        return "", "", ""  # Retorna valores vacíos para evitar errores

    # Convertir el nombre del idioma al código correspondiente
    idioma_dict = dict(IDIOMAS_DISPONIBLES)
    idioma = idioma_dict.get(idioma_nombre, None)

    if not idioma:
        gr.Warning("⚠️ Por favor, selecciona un idioma antes de continuar.")
        return "", "", ""  # Retorna valores vacíos para evitar errores

    # Cargar el modelo solo si es diferente al actual
    model = whisper.load_model(model_name)

    archivos_resultantes = []
    
    for audio_path in audio_paths:
        # Ejecuta la transcripción con Whisper
        result = model.transcribe(audio_path, language=idioma)

        texto_transcrito = result.get("text", "")

        # Si el usuario elige traducir, Whisper generará la traducción al inglés
        if traducir:
            result_traducido = model.transcribe(audio_path, task="translate", language="en")
            texto_transcrito = result_traducido.get("text", "")

        # Crear los nombres de los archivos
        base_name = os.path.splitext(os.path.basename(audio_path))[0]
        txt_file = f"{base_name}.txt"
        srt_file = f"{base_name}.srt"

        # Guardar transcripción en .txt
        with open(txt_file, "w", encoding="utf-8") as f:
            f.write(texto_transcrito)

        # Guardar transcripción en .srt
        with open(srt_file, "w", encoding="utf-8") as f:
            for i, segment in enumerate(result["segments"], start=1):
                start = segment["start"]
                end = segment["end"]
                text = segment["text"]

                f.write(f"{i}\n")
                f.write(f"{start:.3f} --> {end:.3f}\n")
                f.write(f"{text}\n\n")

        archivos_resultantes.append((texto_transcrito, txt_file, srt_file))

    if len(archivos_resultantes) == 1:
        return archivos_resultantes[0]  # Devuelve un solo conjunto de archivos
    else:
        transcripciones, txt_files, srt_files = zip(*archivos_resultantes)  # Separa en listas
        return list(transcripciones), list(txt_files), list(srt_files)  # Devuelve listas en lugar de una tupla

# Interfaz de Gradio con selección de modelo, idioma y opción de traducción
iface = gr.Interface(
    fn=transcribir,
    inputs=[
        gr.Files(type="filepath", label="Sube uno o más archivos"),
        gr.Dropdown(MODEL_OPTIONS, value=None, label="Selecciona el modelo de Whisper", interactive=True),
        gr.Dropdown(
            choices=[nombre for nombre, _ in IDIOMAS_DISPONIBLES],  # Muestra nombres en el menú
            value=None,
            label="Selecciona el idioma del audio",
            interactive=True
        ),
        gr.Checkbox(label="Traducir al inglés", value=False)
    ],
    outputs=[
        gr.Textbox(label="Transcripción"),
        gr.File(label="Descargar .txt"),
        gr.File(label="Descargar .srt")
    ],
    title="Transcriptor de Audio con Whisper",
    description="Sube uno o más archivos de audio y obtén la transcripción con opciones avanzadas."
)

iface.launch(share=True, pwa=True)