import gradio as gr
from transformers import pipeline
from huggingface_hub import InferenceClient
import os
import json

# Initialize Hugging Face Inference Client
client = InferenceClient(api_key="your_huggingface_api_key_here")

# 🎯 Study-only question filter
def is_study_related(question):
    educational_keywords = [
        "math", "science", "ict", "english", "chemistry", "physics", "biology",
        "grammar", "essay", "study", "lesson", "equation", "formula", "computer",
        "programming", "AI", "machine learning", "technology", "education",
        "subject", "exam", "revision", "teacher", "learning", "school", "topic"
    ]

    for word in educational_keywords:
        if word.lower() in question.lower():
            return True
    return False

# Memory save/load
def save_memory(history):
    with open("chat_memory.json", "w") as f:
        json.dump(history, f)

def load_memory():
    if os.path.exists("chat_memory.json"):
        with open("chat_memory.json", "r") as f:
            return json.load(f)
    return []

# Chat logic
def chat_with_model(message, history):
    if not message:
        return history, history

    # 🚫 Block unnecessary/off-topic questions
    if not is_study_related(message):
        reply = "🚫 I'm sorry, but I can only answer study-related questions. Let's focus on learning!"
        history.append((message, reply))
        save_memory(history)
        return history, history

    # Append user message to history
    history.append((message, ""))
    save_memory(history)

    # Generate AI response
    response = client.text_generation(
        model="mistralai/Mixtral-8x7B-Instruct-v0.1",
        prompt=message,
        max_new_tokens=300,
        temperature=0.7
    )

    reply = response
    history[-1] = (message, reply)
    save_memory(history)

    return history, history

# Load existing memory
memory = load_memory()

# Interface
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("## 🤖 EduAI — Where Curiosity Meets Knowledge")
    chatbot = gr.Chatbot(label="EduAI Learning Assistant", value=memory)
    msg = gr.Textbox(label="Ask EduAI a study question...")
    clear = gr.Button("Clear Chat")

    msg.submit(chat_with_model, [msg, chatbot], [chatbot, chatbot])
    clear.click(lambda: [], None, chatbot)

# Launch app
demo.launch()