import os, shlex, subprocess, tempfile, traceback, glob, gc, shutil
import time
import psutil
import humanize
import gradio as gr

# Les imports lourds (torch, nemo) sont déplacés dans les fonctions
# pour assurer un démarrage léger de l'Espace.

# Configuration
DEVICE = "cpu" # Forcé CPU pour stabilité au démarrage, sera vérifié dynamiquement
MODEL_CACHE = {}

def get_model(name, repo, arch_type):
    """Téléchargement et chargement différé d'un modèle NeMo"""
    try:
        # Patch pour éviter l'erreur de compatibilité ConfidenceConfig avec tdt_include_duration
        try:
            import nemo.collections.asr.parts.utils.asr_confidence_utils as asr_confidence_utils
            import inspect
            original_init = asr_confidence_utils.ConfidenceConfig.__init__
            def patched_init(self, *args, **kwargs):
                try:
                    valid_keys = set(inspect.signature(asr_confidence_utils.ConfidenceConfig).parameters.keys())
                    filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_keys}
                except Exception:
                    filtered_kwargs = kwargs
                original_init(self, *args, **filtered_kwargs)
            asr_confidence_utils.ConfidenceConfig.__init__ = patched_init
            print("🩹 Successfully patched NeMo ConfidenceConfig.__init__")
        except Exception as e:
            print(f"⚠️ Could not patch NeMo ConfidenceConfig: {e}")

        import torch
        from huggingface_hub import snapshot_download
        from nemo.collections.asr.models import EncDecCTCModel, EncDecRNNTModel, EncDecHybridRNNTCTCBPEModel
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        token = os.environ.get("HF_TOKEN")
        
        print(f"📥 Folder download: {repo}...")
        folder = snapshot_download(repo, local_dir_use_symlinks=False, token=token)
        nemo_file = glob.glob(os.path.join(folder, "*.nemo"))[0]
        
        print(f"📥 Loading model {name} on {device}...")
        if arch_type == "ctc":
            model = EncDecCTCModel.restore_from(nemo_file, map_location=torch.device(device))
        elif "soloni" in repo.lower():
            model = EncDecHybridRNNTCTCBPEModel.restore_from(nemo_file, map_location=torch.device(device))
        else:
            model = EncDecRNNTModel.restore_from(nemo_file, map_location=torch.device(device))
            
        model.eval()
        return model
    except Exception as e:
        print(f"❌ Erreur lors du chargement du modèle {name}:")
        print(traceback.format_exc())
        return None

MODELS = {
    "Soloba V3 (CTC)":           ("RobotsMali/soloba-ctc-0.6b-v3", "ctc"),
    "Soloni V3 (TDT-CTC)":       ("RobotsMali/soloni-114m-tdt-ctc-v3", "hybrid"),
}

def pipeline(audio_in, model_name, progress=gr.Progress()):
    if not audio_in:
        return "⚠️ Erreur", None, "Veuillez fournir un fichier audio."
    
    try:
        import torch
        import psutil
        
        # Nettoyage mémoire préventif
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        repo, arch_type = MODELS[model_name]
        
        # Chargement à la demande
        progress(0.2, desc="Chargement du modèle...")
        model = get_model(model_name, repo, arch_type)
        
        if model is None:
            return "❌ Erreur", None, f"Impossible de charger le modèle {model_name}. Vérifiez les logs."

        progress(0.5, desc="Transcription en cours...")
        transcription = model.transcribe([audio_in])[0]
        
        # Nettoyage post-exécution
        del model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
        return "✅ Succès", None, transcription
    except Exception as e:
        return "❌ Erreur", None, str(e)

# UI
with gr.Blocks(title="RobotsMali ASR", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎙️ RobotsMali ASR Demo")
    gr.Markdown("Transcription automatique de la parole pour le Bambara et les langues du Mali.")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio")
            model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Soloni V3 (TDT-CTC)", label="Modèle")
            submit_btn = gr.Button("Transcrire", variant="primary")
            
        with gr.Column():
            status_out = gr.Textbox(label="Statut")
            text_output = gr.Textbox(label="Transcription", lines=10)

    submit_btn.click(
        fn=pipeline,
        inputs=[audio_input, model_dropdown],
        outputs=[status_out, gr.State(), text_output]
    )

if __name__ == "__main__":
    print("🚀 Démarrage de l'interface Gradio...")
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)