import gradio as gr
import torch
import os
from huggingface_hub import hf_hub_download
import sys
import soundfile as sf
import numpy as np
import logging
import tempfile
import spaces

# Configuration du logger
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

#Chemin local de téléchargement des fichiers (il faut s'assurer que le dossier existe)
LOCAL_DOWNLOAD_PATH = os.path.dirname(os.path.abspath(__file__)) # Utiliser le chemin du script
# Télécharger le script d'inférence
repo_id = "dofbi/galsenai-xtts-v2-wolof-inference"
inference_file = hf_hub_download(repo_id=repo_id, filename="inference.py", local_dir=LOCAL_DOWNLOAD_PATH)

# Ajouter le dossier au chemin de recherche
sys.path.insert(0, LOCAL_DOWNLOAD_PATH)

# Importer la classe à partir du script d'inférence téléchargé
from inference import WolofXTTSInference

# Initialiser le modèle une seule fois
tts_model = WolofXTTSInference()

@spaces.GPU(duration=120)
def tts(text: str, audio_reference: tuple[int, np.ndarray]) -> tuple[int, np.ndarray] | str:
    """
    Synthétise de la parole à partir d'un texte en utilisant un audio de référence.

    Args:
        text (str): Le texte à synthétiser.
        audio_reference (tuple[int, np.ndarray]): Un tuple contenant le taux d'échantillonnage et les données audio de référence.

    Returns:
         tuple[int, np.ndarray] | str: un tuple contenant le taux d'échantillonnage et les données audio synthétisées, ou un message d'erreur.
    """
    logging.debug(f"tts function called with text: {text} and audio_reference: {audio_reference}")

    if not text or audio_reference is None:
        logging.debug("Text or audio reference is missing")
        return "Veuillez entrer un texte et fournir un audio de référence."

    try:
        sample_rate, audio_data = audio_reference

        # Créer un fichier temporaire pour l'audio de référence
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as temp_audio_file:
            sf.write(temp_audio_file.name, audio_data, sample_rate)
            logging.debug(f"Audio reference saved to {temp_audio_file.name}")

            # Utiliser la méthode generate_audio de la nouvelle classe
            audio_output, output_sample_rate = tts_model.generate_audio(
                text, 
                reference_audio=temp_audio_file.name
            )

            logging.debug(f"Audio generated with sample rate: {output_sample_rate}")
            return (output_sample_rate, audio_output)

    except Exception as e:
        logging.error(f"Error during audio generation: {e}")
        return f"Une erreur s'est produite lors de la génération audio: {e}"

if __name__ == "__main__":
    demo = gr.Interface(
        fn=tts,
        inputs=[
            gr.Textbox(label="Text to synthesize"),
            gr.Audio(type="numpy", label="Reference audio")
        ],
        outputs=gr.Audio(label="Synthesized audio"),
    )

    demo.launch()