"""
Répétiteur Vocal — POC (Bac tchadien, Maths série D, français)
Interface : réplique WhatsApp (le produit final vivra sur WhatsApp).
Pipeline : Vocale élève → ASR (faster-whisper) → LLM (Qwen2.5-1.5B local, GGUF)
           → TTS (MMS-TTS français) → Vocale réponse

100% modèles open-source, 100% local : AUCUN token ni secret requis.
Conçu pour un Space HF gratuit (CPU 2 vCPU, 16 Go RAM).
"""

import os
import re
import tempfile
from datetime import datetime

import gradio as gr
import numpy as np
import scipy.io.wavfile as wavfile
import torch
from faster_whisper import WhisperModel
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from transformers import AutoTokenizer, VitsModel

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
LLM_REPO = "Qwen/Qwen2.5-1.5B-Instruct-GGUF"
LLM_FILE = "qwen2.5-1.5b-instruct-q4_k_m.gguf"
ASR_MODEL_SIZE = "small"
TTS_MODEL = "facebook/mms-tts-fra"

SYSTEM_PROMPT = """Tu es un répétiteur de mathématiques pour des élèves de \
Terminale D au Tchad qui préparent le baccalauréat. Tu expliques comme un \
grand frère patient, en français simple et clair.

Règles impératives :
1. Ta réponse sera LUE À VOIX HAUTE : écris tout en toutes lettres. \
Jamais de symboles mathématiques. Écris "x au carré" et non "x²", \
"racine de deux" et non "√2", "un demi" et non "1/2", "moins trois" et non "-3".
2. Pas de listes, pas de titres, pas de formatage : uniquement des phrases \
courtes qui s'enchaînent naturellement à l'oral.
3. Maximum 150 mots. Va à l'essentiel : l'idée clé, puis un petit exemple.
4. Utilise des exemples concrets de la vie au Tchad quand c'est possible \
(marché, francs CFA, distances entre villes).
5. Termine par une question courte pour vérifier que l'élève a compris.
6. Si la question ne concerne pas les études, ramène gentiment l'élève \
vers ses révisions."""

# ---------------------------------------------------------------------------
# Chargement des modèles locaux (une seule fois au démarrage du Space)
# ---------------------------------------------------------------------------
print("Chargement ASR (faster-whisper)...")
asr_model = WhisperModel(ASR_MODEL_SIZE, device="cpu", compute_type="int8")

print("Chargement TTS (MMS-TTS français)...")
tts_model = VitsModel.from_pretrained(TTS_MODEL)
tts_tokenizer = AutoTokenizer.from_pretrained(TTS_MODEL)
tts_model.eval()

print("Chargement LLM local (Qwen2.5-1.5B GGUF, Q4)...")
llm_path = hf_hub_download(repo_id=LLM_REPO, filename=LLM_FILE)
llm = Llama(model_path=llm_path, n_ctx=2048, n_threads=2, verbose=False)
print("Modèles prêts.")

# ---------------------------------------------------------------------------
# Briques du pipeline
# ---------------------------------------------------------------------------
def transcrire(audio_path: str) -> str:
    segments, _ = asr_model.transcribe(audio_path, language="fr", beam_size=5)
    return " ".join(seg.text.strip() for seg in segments).strip()


def repondre(question: str, memoire: list) -> str:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for q, r in memoire[-3:]:
        messages.append({"role": "user", "content": q})
        messages.append({"role": "assistant", "content": r})
    messages.append({"role": "user", "content": question})
    reponse = llm.create_chat_completion(
        messages=messages, max_tokens=250, temperature=0.4
    )
    return reponse["choices"][0]["message"]["content"].strip()


def nettoyer_pour_tts(texte: str) -> str:
    texte = re.sub(r"[*_#`>\[\]()]", " ", texte)
    texte = re.sub(r"\s+", " ", texte)
    return texte.strip()


def synthetiser(texte: str) -> str:
    texte = nettoyer_pour_tts(texte)
    inputs = tts_tokenizer(texte, return_tensors="pt")
    with torch.no_grad():
        waveform = tts_model(**inputs).waveform
    audio = waveform.squeeze().cpu().numpy()
    audio = (audio / np.max(np.abs(audio)) * 32767).astype(np.int16)
    out_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    wavfile.write(out_path, tts_model.config.sampling_rate, audio)
    return out_path


def _heure() -> str:
    return datetime.now().strftime("%H:%M")


# ---------------------------------------------------------------------------
# Logique de l'interface (format "messages" du Chatbot Gradio)
# ---------------------------------------------------------------------------
def traiter(audio_path, question_texte, chat, memoire):
    chat = chat or []
    memoire = memoire or []

    if audio_path:
        question = transcrire(audio_path)
        etiquette = f"🎤 {question}" if question else ""
    elif question_texte and question_texte.strip():
        question = question_texte.strip()
        etiquette = question
    else:
        return chat, memoire, None, ""

    if not question:
        chat.append({"role": "assistant",
                     "content": "Je n'ai pas bien entendu 🙏 Réessaie en parlant plus fort."})
        return chat, memoire, None, ""

    chat.append({"role": "user", "content": etiquette})

    try:
        explication = repondre(question, memoire)
    except Exception as e:
        chat.append({"role": "assistant",
                     "content": f"⚠️ Erreur LLM : {e}"})
        return chat, memoire, None, ""

    # Note vocale du répétiteur, puis transcription texte (comme sur WhatsApp)
    try:
        audio_reponse = synthetiser(explication)
        chat.append({"role": "assistant",
                     "content": {"path": audio_reponse, "mime_type": "audio/wav"}})
    except Exception:
        pass
    chat.append({"role": "assistant", "content": explication})

    memoire.append((question, explication))
    return chat, memoire, None, ""


# ---------------------------------------------------------------------------
# Interface — réplique WhatsApp
# ---------------------------------------------------------------------------
CSS = """
/* ---- cadre téléphone ---- */
.gradio-container {
    max-width: 460px !important;
    margin: 0 auto !important;
    background: #0b141a !important;
    padding: 0 !important;
}
footer {display: none !important;}

/* ---- en-tête WhatsApp ---- */
#wa-header {
    background: #075E54;
    color: #fff;
    padding: 10px 14px;
    border-radius: 0;
    margin: 0;
}
#wa-header * {color: #fff !important;}

/* ---- zone de chat : papier peint WhatsApp ---- */
#wa-chat {
    background-color: #ECE5DD !important;
    background-image: radial-gradient(#d9d2c9 0.75px, transparent 0.75px);
    background-size: 18px 18px;
    border: none !important;
    border-radius: 0 !important;
}
#wa-chat .placeholder, #wa-chat .panel {background: transparent !important;}

/* ---- bulles ---- */
#wa-chat .message {
    border: none !important;
    border-radius: 8px !important;
    box-shadow: 0 1px 0.5px rgba(0,0,0,0.13) !important;
    font-size: 14.5px !important;
    max-width: 82% !important;
}
#wa-chat .message.user, #wa-chat .user {
    background: #DCF8C6 !important;
    color: #111b21 !important;
}
#wa-chat .message.bot, #wa-chat .bot {
    background: #ffffff !important;
    color: #111b21 !important;
}
#wa-chat .message audio {width: 230px; height: 40px;}
#wa-chat .avatar-container {display: none !important;}

/* ---- barre de saisie ---- */
#wa-input {
    background: #F0F2F5 !important;
    padding: 6px 8px !important;
    margin: 0 !important;
    border-radius: 0 !important;
}
#wa-text textarea {
    border-radius: 20px !important;
    background: #fff !important;
    border: none !important;
}
#wa-send {
    background: #00A884 !important;
    color: #fff !important;
    border-radius: 50% !important;
    min-width: 46px !important;
    max-width: 46px !important;
    height: 46px !important;
    font-size: 20px !important;
}
#wa-mic {border: none !important; background: transparent !important;}
#wa-mic .wrap {border: none !important;}
"""

EN_TETE = """
<div style="display:flex;align-items:center;gap:12px;">
  <div style="width:40px;height:40px;border-radius:50%;background:#128C7E;
              display:flex;align-items:center;justify-content:center;font-size:20px;">🎓</div>
  <div style="flex:1;">
    <div style="font-weight:600;font-size:16px;">Répétiteur Maths D</div>
    <div style="font-size:12px;opacity:0.85;">en ligne</div>
  </div>
  <div style="font-size:18px;opacity:0.9;">📹&nbsp;&nbsp;📞&nbsp;&nbsp;⋮</div>
</div>
"""

with gr.Blocks(css=CSS, title="Répétiteur Vocal", theme=gr.themes.Base()) as demo:
    gr.HTML(EN_TETE, elem_id="wa-header")

    chatbot = gr.Chatbot(
        value=[{"role": "assistant",
                "content": "Salut ! 👋 Je suis ton répétiteur de maths Terminale D. "
                           "Envoie-moi ta question en note vocale 🎤 ou par écrit."}],
        type="messages",
        elem_id="wa-chat",
        height=520,
        show_label=False,
    )
    memoire = gr.State([])

    with gr.Row(elem_id="wa-input"):
        micro = gr.Audio(
            sources=["microphone"], type="filepath", show_label=False,
            elem_id="wa-mic", scale=2, waveform_options={"show_recording_waveform": False},
        )
        texte = gr.Textbox(
            show_label=False, placeholder="Message", elem_id="wa-text", scale=3,
        )
        envoyer = gr.Button("➤", elem_id="wa-send", scale=0)

    envoyer.click(
        traiter,
        inputs=[micro, texte, chatbot, memoire],
        outputs=[chatbot, memoire, micro, texte],
    )
    texte.submit(
        traiter,
        inputs=[micro, texte, chatbot, memoire],
        outputs=[chatbot, memoire, micro, texte],
    )
    # Envoi automatique dès la fin de l'enregistrement vocal (comme WhatsApp)
    micro.stop_recording(
        traiter,
        inputs=[micro, texte, chatbot, memoire],
        outputs=[chatbot, memoire, micro, texte],
    )

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