import gradio as gr
from gradio.components import Text
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
import torch
import torchaudio
import librosa
import numpy as np


if torch.backends.mps.is_available():
    device = torch.device("mps")
elif torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

# Pipeline de transcription, c’est-à-dire qui permet de transcrire l’audio en texte français
transcriber = pipeline(
    "automatic-speech-recognition", 
    model="bhuang/asr-wav2vec2-french", 
    device=device
)

#Pipeline de traduction qui va permettre de traduire l’audio transcrit en anglais
translator_model = "Helsinki-NLP/opus-mt-fr-en"
translator = pipeline(
    "translation_fr_to_en",
    model=translator_model,
    tokenizer=translator_model,
    device=device
)
# Pre-traitement necessaire 
def preprocess_audio(filepath):
    """
    Charge l'audio, le normalise, supprime le bruit et le rééchantillonne.
    Retourne le chemin du fichier temporaire prétraité.
    """
    # Chargement et rééchantillonnage
    audio, sr = librosa.load(filepath, sr=16000)
    
    # Normalisation
    audio = librosa.util.normalize(audio)
    
    # Suppression du bruit simple par seuillage (option basique)
    noise_threshold = 0.01
    audio = np.where(np.abs(audio) < noise_threshold, 0, audio)
    
    # Sauvegarde du fichier temporaire
    temp_path = "preprocessed_audio.wav"
    torchaudio.save(temp_path, torch.tensor(audio).unsqueeze(0), sample_rate=16000)
    
    return temp_path

def pipeline_transcription_traduction(audio):
    """
    Prend un fichier audio, le prétraite, le transcrit (FR), puis le traduit (EN).
    """
    try:
        # Étape 1 : Prétraitement audio
        audio_path = preprocess_audio(audio)
        
        # Étape 2 : Transcription
        transcription = transcriber(audio_path)["text"]
        
        # Étape 3 : Traduction
        traduction = translator(transcription)[0]["translation_text"]
        
        return transcription, traduction

    except Exception as e:
        return f"Erreur : {str(e)}", "Erreur dans la traduction"

# Interface Gradio
demo = gr.Interface(
    title="Transcription & Traduction Audio FR → EN",
    description=f"Traitement sur : {device}",
    fn=pipeline_transcription_traduction,
    inputs=[gr.Audio(sources=["microphone"], type="filepath", label="Audio")],
    outputs=[
        Text(label="Transcription (Français)"), 
        Text(label="Traduction (Anglais)")
    ]
)

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