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)

        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
        )

        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
"""

# JavaScript mejorado para limpiar los datos de autenticación
clear_auth_js = """
function clearAuthData() {
    localStorage.removeItem('gradio_auth_token');
    localStorage.removeItem('gradio_auth_expiration');
    sessionStorage.removeItem('gradio_auth_token');
    sessionStorage.removeItem('gradio_auth_expiration');
    document.cookie = 'gradio_auth_token=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;';
    document.cookie = 'gradio_auth_expiration=; expires=Thu, 01 Jan 1970 00:00:00 UTC; path=/;';
}

window.addEventListener('beforeunload', clearAuthData);

function logout() {
    clearAuthData();
    window.location.reload();
}
"""

# CSS personalizado
custom_css = """
#image-container img {
    display: block;
    margin-left: auto;
    margin-right: auto;
    max-width: 256px;
    height: auto;
}
.logout-button {
    position: fixed;
    top: 10px;
    right: 10px;
    z-index: 1000;
    padding: 8px 16px;
    background-color: #f44336;
    color: white;
    border: none;
    border-radius: 4px;
    cursor: pointer;
}
.logout-button:hover {
    background-color: #d32f2f;
}
.login-container {
    background-color: white;
    padding: 2rem;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    text-align: center;
    max-width: 400px;
    width: 100%;
}
.login-container h1 {
    margin-bottom: 1rem;
    color: #4a4a4a;
}
.login-container input {
    width: 100%;
    padding: 0.5rem;
    margin-bottom: 1rem;
    border: 1px solid #ddd;
    border-radius: 4px;
}
.login-container button {
    width: 100%;
    padding: 0.5rem;
    background-color: #3498db;
    color: white;
    border: none;
    border-radius: 4px;
    cursor: pointer;
}
.login-container button:hover {
    background-color: #2980b9;
}
"""

# Modificar la parte del formulario de inicio de sesión
def custom_auth(username, password):
    if (username, password) in [("Pedro Labattaglia", "PL2024"), ("Invitado", "PLTTS2024")]:
        return True
    return False

iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Texto a sintetizar", placeholder="Escribe aquí el texto que quieres convertir a voz..."),
        gr.Dropdown(label="Idioma", choices=supported_languages),
        gr.Dropdown(label="Audio de referencia", choices=reference_audios)
    ],
    outputs=[
        gr.Audio(label="Audio generado"),
        gr.Textbox(label="Métricas")
    ],
    title="Sintetizador de voz de Pedro Labattaglia",
    description=description,
    theme=theme,
    css=custom_css,
    allow_flagging="never"
)

# Crear una nueva interfaz para el inicio de sesión
login_iface = gr.Interface(
    fn=custom_auth,
    inputs=[
        gr.Textbox(label="Usuario", placeholder="Ingrese su nombre de usuario"),
        gr.Textbox(label="Contraseña", type="password", placeholder="Ingrese su contraseña")
    ],
    outputs=gr.Textbox(visible=False),
    title="Bienvenido al sintetizador de voz de Pedro Labattaglia",
    description="Por favor, introduzca sus credenciales para acceder.",
    theme=theme,
    css=custom_css
)

# Combinar las interfaces
demo = gr.TabbedInterface([login_iface, iface], ["Login", "Sintetizador"])

if __name__ == "__main__":
    demo.launch()