import gradio as gr
import torch
import torchaudio
import librosa
import os
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq
from huggingface_hub import login
import logging

# Configuration
MODEL_NAME = "Ronaldodev/speech-to-text-fongbe"
HF_TOKEN = os.environ.get("HF_TOKEN")

# Variables globales
model = None
processor = None

# Configuration logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def load_model():
    """Charger le modèle privé au démarrage"""
    global model, processor

    try:
        logger.info("🔄 Chargement du modèle privé...")

        if not HF_TOKEN:
            raise ValueError("HF_TOKEN non configuré dans les secrets")

        # Login avec token privé
        login(token=HF_TOKEN)
        logger.info("✅ Authentification HF réussie")

        # Charger le modèle et processeur
        model = AutoModelForSpeechSeq2Seq.from_pretrained(MODEL_NAME)
        processor = AutoProcessor.from_pretrained(MODEL_NAME)

        logger.info("✅ Modèle chargé avec succès!")
        return True

    except Exception as e:
        logger.error(f"❌ Erreur chargement: {e}")
        return False


def transcribe(audio):
    """Fonction principale de transcription"""

    # Vérifier si le modèle est chargé
    if model is None or processor is None:
        return "❌ Erreur: Modèle non chargé. Vérifiez les logs."

    # Vérifier si un audio est fourni
    if audio is None:
        return "❌ Aucun fichier audio fourni"

    try:
        logger.info(f"🎵 Traitement audio: {audio}")

        # Charger l'audio avec fallback
        try:
            waveform, sample_rate = torchaudio.load(audio)
            logger.info(f"✅ Audio chargé avec torchaudio: {sample_rate}Hz")
        except Exception as e:
            logger.warning(f"⚠️ Torchaudio échoué, essai librosa: {e}")
            waveform, sample_rate = librosa.load(audio, sr=None)
            waveform = torch.tensor(waveform).unsqueeze(0)
            logger.info(f"✅ Audio chargé avec librosa: {sample_rate}Hz")

        # Conversion mono si nécessaire
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
            logger.info("🔄 Conversion stéréo → mono")

        # Resampling à 16kHz si nécessaire
        if sample_rate != 16000:
            logger.info(f"🔄 Resampling {sample_rate}Hz → 16000Hz")
            resampler = torchaudio.transforms.Resample(sample_rate, 16000)
            waveform = resampler(waveform)

        # Préparation des inputs
        inputs = processor(
            waveform.squeeze(),
            sampling_rate=16000,
            return_tensors="pt"
        )

        # Génération de la transcription
        logger.info("🔄 Génération de la transcription...")
        with torch.no_grad():
            result = model.generate(
                **inputs,
                max_length=500,
                do_sample=False,
                num_beams=1
            )

        # Décodage
        transcription = processor.batch_decode(result, skip_special_tokens=True)[0]

        logger.info(f"✅ Transcription réussie: '{transcription}'")
        return transcription.strip()

    except Exception as e:
        error_msg = f"❌ Erreur de transcription: {str(e)}"
        logger.error(error_msg)
        return error_msg


# Charger le modèle au démarrage
print("🚀 DÉMARRAGE API STT FONGBÉ - RONALDODEV")
print("=" * 50)

if load_model():
    print("✅ Modèle chargé - Interface prête!")
    model_status = "✅ Modèle chargé et prêt"
else:
    print("❌ Erreur de chargement du modèle")
    model_status = "❌ Erreur de chargement"

# Interface Gradio simple
demo = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio(
        sources=["upload", "microphone"],
        type="filepath",
        label="🎤 Uploadez un fichier ou enregistrez directement"
    ),
    outputs=gr.Textbox(
        label="📝 Transcription en Fongbé",
        placeholder="La transcription apparaîtra ici...",
        lines=3
    ),
    title="🎤 API STT Fongbé - Ronaldodev",
    description=f"""
    **Reconnaissance vocale pour la langue Fongbé**

    Uploadez un fichier audio (WAV, MP3, M4A) ou enregistrez directement avec votre microphone.

    **Statut:** {model_status}

    **Modèle:** `{MODEL_NAME}`
    """,
    article="""
    ## 🔌 API pour développeurs

    Cette interface expose automatiquement une API REST :

    **Endpoint:** `POST /api/predict`

    **Exemple d'utilisation:**
    ```python
    import requests

    response = requests.post(
        "https://ronaldodev-stt-fongbe.hf.space/api/predict",
        json={"data": [audio_file_path]}
    )

    transcription = response.json()["data"][0]
    ```

    **Pour Flutter:** Utilisez MultipartRequest avec l'endpoint ci-dessus.
    """,
    examples=[
        # Vous pouvez ajouter des fichiers d'exemple si vous en avez
    ],
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

# Lancement de l'interface
demo.launch()