import os
import re
import time
import sys
import subprocess
import scipy.io.wavfile as wavfile
import torch
import torchaudio
import gradio as gr
from TTS.api import TTS
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from TTS.utils.generic_utils import get_user_data_dir
from huggingface_hub import hf_hub_download

# Configuración inicial
os.environ["COQUI_TOS_AGREED"] = "1"

def check_and_install(package):
    try:
        __import__(package)
    except ImportError:
        print(f"{package} no está instalado. Instalando...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])

print("Descargando y configurando el modelo...")
repo_id = "Blakus/Pedro_Lab_XTTS"
local_dir = os.path.join(get_user_data_dir("tts"), "tts_models--multilingual--multi-dataset--xtts_v2")
os.makedirs(local_dir, exist_ok=True)
files_to_download = ["config.json", "model.pth", "vocab.json"]

for file_name in files_to_download:
    print(f"Descargando {file_name} de {repo_id}")
    hf_hub_download(repo_id=repo_id, filename=file_name, local_dir=local_dir)

config_path = os.path.join(local_dir, "config.json")
checkpoint_path = os.path.join(local_dir, "model.pth")
vocab_path = os.path.join(local_dir, "vocab.json")

config = XttsConfig()
config.load_json(config_path)

model = Xtts.init_from_config(config)
model.load_checkpoint(config, checkpoint_path=checkpoint_path, vocab_path=vocab_path, eval=True, use_deepspeed=True)

model.cuda()

print("Modelo cargado en GPU")

def predict(prompt, language, reference_audio):
    try:
        if len(prompt) < 2 or len(prompt) > 600:
            return None, "El texto debe tener entre 2 y 600 caracteres."

        # Obtener los parámetros de la configuración JSON
        temperature = config.model_args.get("temperature", 0.85)
        length_penalty = config.model_args.get("length_penalty", 1.0)
        repetition_penalty = config.model_args.get("repetition_penalty", 2.0)
        top_k = config.model_args.get("top_k", 50)
        top_p = config.model_args.get("top_p", 0.85)

        # Parámetros agregados
        num_gpt_outputs = config.model_args.get("num_gpt_outputs", 32)
        gpt_cond_len = config.model_args.get("gpt_cond_len", 24)
        gpt_cond_chunk_len = config.model_args.get("gpt_cond_chunk_len", 4)
        max_ref_len = config.model_args.get("max_ref_len", 10)

        gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(
            audio_path=reference_audio
        )

        start_time = time.time()

        out = model.inference(
            prompt,
            language,
            gpt_cond_latent,
            speaker_embedding,
            temperature=temperature,
            length_penalty=length_penalty,
            repetition_penalty=repetition_penalty,
            top_k=top_k,
            top_p=top_p,
            num_gpt_outputs=num_gpt_outputs,       # Número de muestras del modelo
            gpt_cond_len=gpt_cond_len,             # Longitud del condicionamiento
            gpt_cond_chunk_len=gpt_cond_chunk_len, # Tamaño del fragmento de audio
            max_ref_len=max_ref_len                # Máximo de segundos para condicionamiento
        )

        inference_time = time.time() - start_time
        
        output_path = "pedro_labattaglia_TTS.wav"
        # Guardar el audio directamente desde el output del modelo
        wavfile.write(output_path, config.audio["output_sample_rate"], out["wav"])

        audio_length = len(out["wav"]) / config.audio["output_sample_rate"]  # duración del audio en segundos
        real_time_factor = inference_time / audio_length

        metrics_text = f"Tiempo de generación: {inference_time:.2f} segundos\n"
        metrics_text += f"Factor de tiempo real: {real_time_factor:.2f}"

        return output_path, metrics_text

    except Exception as e:
        print(f"Error detallado: {str(e)}")
        return None, f"Error: {str(e)}"

# Configuración de la interfaz de Gradio
supported_languages = ["es", "en"]
reference_audios = [
    "serio.wav",
    "neutral.wav",
    "alegre.wav",
    "neutral_ingles.wav"
]

theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="gray",
).set(
    body_background_fill='*neutral_100',
    body_background_fill_dark='*neutral_900',
)

description = """
# Sintetizador de voz de Pedro Labattaglia 🎙️
Sintetizador de voz con la voz del locutor argentino Pedro Labattaglia. 
## Cómo usarlo:
- Elija el idioma (Español o Inglés)
- Elija un audio de referencia de la lista 
- Escriba el texto que desea sintetizar
- Presione generar voz
"""

# Interfaz de Gradio
with gr.Blocks(theme=theme) as demo:
    gr.Markdown(description)

    # Fila para centrar la imagen
    with gr.Row():
        with gr.Column(equal_height=True):
            gr.Image(
                "https://www.labattaglia.com.ar/images/about_me_pic2.jpg", 
                label="", 
                show_label=False,
                container=False,
                elem_id="image-container"
            )

    # Fila para seleccionar idioma, referencia y generar voz
    with gr.Row():
        with gr.Column(scale=2):
            language_selector = gr.Dropdown(label="Idioma", choices=supported_languages)
            reference_audio = gr.Dropdown(label="Audio de referencia", choices=reference_audios)
            input_text = gr.Textbox(label="Texto a sintetizar", placeholder="Escribe aquí el texto que quieres convertir a voz...")
            generate_button = gr.Button("Generar voz", variant="primary")

        with gr.Column(scale=1):
            generated_audio = gr.Audio(label="Audio generado", interactive=False)
            metrics_output = gr.Textbox(label="Métricas", value="Tiempo de generación: -- segundos\nFactor de tiempo real: --")

    # Configuración del botón para generar voz
    generate_button.click(
        predict,
        inputs=[input_text, language_selector, reference_audio],
        outputs=[generated_audio, metrics_output]
    )

# Estilos CSS personalizados
demo.css = """
#image-container img {
    display: block;
    margin-left: auto;
    margin-right: auto;
    max-width: 256px;
    height: auto;
}
"""

if __name__ == "__main__":
    demo.launch(auth=[("Pedro Labattaglia", "PL2024"), ("Invitado", "PLTTS2024")])