import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2

# Initialisation du modèle HF中国镜像站
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Messages système pour guider le modèle
SYSTEM_PROMPT = {
    "fr": "Tu es un assistant pédagogique qui aide les professeurs à créer des cours et analyser des documents PDF.",
    "en": "You are an educational assistant helping teachers create courses and analyze PDF documents."
}

# 📄 Fonction pour lire et extraire le texte d'un PDF
def extract_text_from_pdf(pdf_file):
    text = ""
    with pdf_file as f:
        reader = PyPDF2.PdfReader(f)
        for page in reader.pages:
            text += page.extract_text() + "\n"
    return text

# 🧠 Fonction du chatbot avec gestion de l'historique + PDF RAG
def generate_response(subject, history, lang, pdf_file, max_tokens, temperature, top_p):
    system_message = SYSTEM_PROMPT.get(lang, SYSTEM_PROMPT["en"])  # Sélection de la langue

    messages = [{"role": "system", "content": system_message}]
    
    # 🔄 Correction du format de l'historique
    for message in history:
        if isinstance(message, dict) and "role" in message and "content" in message:
            messages.append(message)

    # 📄 Ajouter le contenu du PDF s'il y en a un
    if pdf_file is not None:
        pdf_text = extract_text_from_pdf(pdf_file)
        messages.append({"role": "user", "content": f"Voici un document PDF pertinent : {pdf_text[:1000]}..."})  # On limite à 1000 caractères pour éviter la surcharge

    # Ajouter la demande de l'utilisateur
    messages.append({"role": "user", "content": f"Crée un cours sur : {subject}"})

    # 🔥 Génération de la réponse en streaming
    response = ""
    for message in client.chat_completion(
        messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# 🎨 Interface utilisateur Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎓 Teacher Assistant Chatbot avec PDF RAG")

    with gr.Row():
        subject_input = gr.Textbox(label="📌 Sujet du cours", placeholder="Ex: Apprentissage automatique")
        lang_select = gr.Dropdown(choices=["fr", "en"], value="fr", label="🌍 Langue")

    pdf_upload = gr.File(label="📄 Télécharger un PDF (optionnel)", type="file")

    chat = gr.Chatbot(type="messages")  # ✅ Correction : Format messages OK

    with gr.Row():
        max_tokens = gr.Slider(minimum=100, maximum=2048, value=512, step=1, label="📝 Max tokens")
        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="🔥 Température")
        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="🎯 Top-p")

    generate_button = gr.Button("🚀 Générer le cours")

    generate_button.click(
        generate_response, 
        inputs=[subject_input, chat, lang_select, pdf_upload, max_tokens, temperature, top_p], 
        outputs=chat
    )

# 🔥 Lancer l'application
if __name__ == "__main__":
    demo.launch()