import gradio as gr
from transformers import pipeline
import scipy
import torchaudio
from speechbrain.pretrained import SepformerSeparation as separator
from shialifube import transliterate
import tempfile
import os

# Chargement des modèles Whisper pour la transcription
model_roman = pipeline("automatic-speech-recognition", model="nairaxo/whisper-shikomori-latin")
model_arabic = pipeline("automatic-speech-recognition", model="nairaxo/whisper-shikomori-arabic")

# Chargement du modèle Text-to-Speech
model_id = "nairaxo/mms-tts-zdj"
synthesiser = pipeline("text-to-speech", model_id, device=0)

# Chargement du modèle d'amélioration audio
model_enh = separator.from_hparams(source="speechbrain/sepformer-wham16k-enhancement", savedir='pretrained_models/sepformer-wham16k-enhancement')

# Fonction de transcription avec sélection du modèle
def transcribe(audio, model_choice):
    try:
        # Vérifier si l'audio est valide
        if audio is None:
            return "Erreur : Aucun fichier audio n'a été fourni."

        # Transcrire l'audio
        if model_choice == "Modèle en alphabet latin":
            transcription = model_roman(audio)["text"]
        else:
            transcription = model_arabic(audio)["text"]

        return transcription

    except Exception as e:
        return f"Erreur lors de la transcription : {str(e)}"

# Fonction de génération et d'amélioration audio
def generate_and_enhance_audio(text, script_choice):
    try:
        # Vérifier si le texte est vide
        if not text.strip():
            return None, None, "Erreur : Le texte d'entrée est vide."

        # Translittérer le texte si l'utilisateur a choisi l'arabe
        if script_choice == "Alphabet arabe":
            text = transliterate(text)  # Translittération de l'arabe en latin

        # Synthétiser la parole (audio original)
        speech = synthesiser(text)
        sampling_rate = speech["sampling_rate"]

        # Créer des fichiers temporaires pour l'audio original et amélioré
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as original_file:
            original_output = original_file.name
            scipy.io.wavfile.write(original_output, rate=sampling_rate, data=speech["audio"][0])

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as enhanced_file:
            enhanced_output = enhanced_file.name
            est_sources = model_enh.separate_file(path=original_output)
            torchaudio.save(enhanced_output, est_sources[:, :, 0].detach().cpu(), sampling_rate)

        # Retourner les fichiers audio au format attendu par Gradio
        return original_output, enhanced_output, None

    except Exception as e:
        return None, None, f"Erreur lors de la génération ou de l'amélioration de l'audio : {str(e)}"

# Interface pour la transcription via microphone
with gr.Blocks() as mf_transcribe:
    gr.Markdown("## Transcription Audio en Shikomori (Microphone)")
    gr.Markdown("<p style='color: #555;'>Sélectionnez une méthode et un modèle pour transcrire l'audio en langue Shikomori. Ce service (en version bêta) prend en charge les transcriptions en alphabet latin et arabe (système Kamar-Eddine). Les modèles ont été entraîné sur la base de données construites à partir d'un algorithme d'alignement forcé. Pour une bonne expérience et afin de mieux transcrire vos audios, assurez-vous de prononcer clairement les mots et d'être dans un environnement ayant peu de bruits.</p>")
    with gr.Row():
        audio_input = gr.Audio(sources=["microphone"], type="filepath", label="🎤 Entrée Audio (Microphone)")
        model_choice = gr.Radio(choices=["Modèle en alphabet latin", "Modèle en alphabet arabe"], label="Sélection du modèle de transcription", value="Modèle en alphabet latin")
    transcription_output = gr.Textbox(label="📄 Transcription en Shikomori", lines=5, max_lines=10)
    transcribe_button = gr.Button("Transcrire")
    transcribe_button.click(fn=transcribe, inputs=[audio_input, model_choice], outputs=transcription_output)

# Interface pour la transcription via fichier audio
with gr.Blocks() as file_transcribe:
    gr.Markdown("## Transcription Audio en Shikomori (Fichier)")
    gr.Markdown("<p style='color: #555;'>Chargez un fichier audio et sélectionnez une méthode et un modèle pour transcrire l'audio en langue Shikomori. Ce service (en version bêta) prend en charge les transcriptions en alphabet latin et arabe (système Kamar-Eddine). Les modèles ont été entraîné sur la base de données construites à partir d'un algorithme d'alignement forcé. Pour une bonne expérience et afin de mieux transcrire vos audios, assurez-vous de prononcer clairement les mots et d'être dans un environnement ayant peu de bruits.</p>")
    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="📂 Entrée Audio (Fichier)")
        model_choice = gr.Radio(choices=["Modèle en alphabet latin", "Modèle en alphabet arabe"], label="Sélection du modèle de transcription", value="Modèle en alphabet latin")
    transcription_output = gr.Textbox(label="📄 Transcription en Shikomori", lines=5, max_lines=10)
    transcribe_button = gr.Button("Transcrire")
    transcribe_button.click(fn=transcribe, inputs=[audio_input, model_choice], outputs=transcription_output)

# Interface pour la synthèse et amélioration audio
with gr.Blocks() as tts_interface:
    gr.Markdown("## 🎙️ Synthèse et amélioration de la parole")
    gr.Markdown("<p style='color: #555;'>Entrez du texte pour générer de la parole en Shikomori. Si le texte est en alphabet arabe, il sera automatiquement translittéré en alphabet latin avant la synthèse. L'audio original et l'audio amélioré seront affichés côte à côte.</p>")
    with gr.Row():
        text_input = gr.Textbox(label="Entrez votre texte", lines=3, placeholder="Écrivez ici...")
        script_choice = gr.Radio(choices=["Alphabet latin", "Alphabet arabe"], label="Sélection du script d'entrée", value="Alphabet latin")
    with gr.Row():
        original_audio = gr.Audio(label="Audio original", type="filepath")
        enhanced_audio = gr.Audio(label="Audio amélioré", type="filepath")
    error_output = gr.Textbox(label="Erreur", visible=False)
    generate_button = gr.Button("Générer l'audio")
    generate_button.click(
        fn=generate_and_enhance_audio,
        inputs=[text_input, script_choice],
        outputs=[original_audio, enhanced_audio, error_output]
    )

# Interface principale avec onglets
with gr.Blocks() as demo:
    gr.TabbedInterface(
        [mf_transcribe, file_transcribe, tts_interface],
        ["🔊 Microphone", "📁 Fichier Audio", "🎙️ Text-to-Speech"]
    )

# Lancement de l'application
demo.launch(share=True, debug=True)