from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import gradio as gr
import re
from pathlib import Path

# ====== CONFIGURATION ======
MODEL_NAME = "microsoft/phi-3-mini-4k-instruct"
SYSTEM_PROMPT_FILE = "infos_medical.txt"  # ton fichier .txt

# ====== CHARGEMENT DU MODÈLE ET TOKENIZER ======
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)
print(f"✅ Modèle '{MODEL_NAME}' chargé avec succès.")

# ====== CHARGEMENT DU CONTEXTE MÉDICAL ======
try:
    context_text = Path(SYSTEM_PROMPT_FILE).read_text(encoding="utf-8").strip()
except FileNotFoundError:
    context_text = "Aucune donnée médicale fournie."
    print("⚠️ Fichier infos_medical.txt introuvable — le chatbot fonctionnera sans contexte.")

# Nettoyage du texte
context_text = context_text.replace("\n", " ")
context_text = re.sub(r"\s+", " ", context_text)
context_text = context_text[:3000]  # limite la longueur pour garder de la mémoire

# ====== INSTRUCTION SYSTÈME ======
system_instruction = (
    "Tu es un assistant médical virtuel. "
    "Tu réponds en français, poliment et de façon factuelle. "
    "Tu fournis uniquement des informations générales de santé. "
    "Tu ne poses pas de diagnostic et ne prescris pas de traitement. "
    "Si la demande est trop spécifique, réponds : "
    "\"Je ne sais pas — consultez un professionnel de santé.\""
)

# ====== CONSTRUCTION DU PROMPT ======
def build_prompt(history, user_message):
    """
    Construit un prompt propre : instruction + contexte + historique + message utilisateur.
    """
    lines = [
        f"System: {system_instruction}",
        "",
        f"Contexte médical de référence : {context_text}",
        "",
    ]
    for user, bot in history:
        lines.append(f"User: {user}")
        if bot:
            lines.append(f"Bot: {bot}")
    lines.append(f"User: {user_message}")
    lines.append("Bot:")
    return "\n".join(lines)

# ====== FONCTION DE RÉPONSE ======
def chat(user_message, history):
    prompt = build_prompt(history, user_message)
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.6,
        top_p=0.9,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Récupération du texte après le dernier "Bot:"
    if "Bot:" in reply:
        reply = reply.split("Bot:")[-1].strip()

    history.append((user_message, reply))
    return history, history

# ====== INTERFACE GRADIO ======
with gr.Blocks(theme=gr.themes.Soft(primary_hue="green")) as demo:
    gr.Markdown("## 🩺 Chatbot Médical — Assistant virtuel de santé")
    gr.Markdown(f"📘 **Contexte chargé :** {len(context_text.split())} mots de référence médicale.")
    chatbot = gr.Chatbot(height=400)
    msg = gr.Textbox(label="Posez une question médicale…")
    clear = gr.Button("🧹 Effacer la conversation")

    msg.submit(chat, [msg, chatbot], [chatbot, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()