import torch
import time
import re
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import gradio as gr

# Load tokenizer and model
model_path = "supp03/E3L-AI-CHATBOT"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto")

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    pad_token_id=tokenizer.eos_token_id
)

# Simple context memory
context_memory = {
    "last_question": None,
    "last_answer": None
}

def generate_response(user_input):
    try:
        if not user_input.strip():
            return "Please enter a question."

        if user_input.strip().lower() == "exit":
            context_memory["last_question"] = None
            context_memory["last_answer"] = None
            return "Session reset. You can ask a new question now."

        refer_prev = any(phrase in user_input.lower() for phrase in [
            "previous question", "previous answer", "elaborate", "detail", "more info", "explain more", "expand on that"
        ])

        context = ""
        if refer_prev and context_memory["last_question"] and context_memory["last_answer"]:
            context = (
                f"Previous Question: {context_memory['last_question']}\n"
                f"Previous Answer: {context_memory['last_answer']}\n\n"
                "Please elaborate on the previous answer with more details, examples, and explanations.\n\n"
            )

        if "steps" in user_input.lower() or "step by step" in user_input.lower():
            style_instruction = "Organize your response into a clearly numbered step-by-step list (1., 2., 3., ...)."
        elif "how" in user_input.lower() or "ways" in user_input.lower():
            style_instruction = "Organize your response into practical bullet points."
        else:
            style_instruction = "Write in a natural, conversational paragraph form."

        generated_prompt = (
            "You are a warm, thoughtful, and emotionally intelligent mentor. "
            "Your job is to provide deeply personalized, practical, and supportive guidance. "
            "Avoid sounding robotic or generic. Respond with empathy, clarity, and human warmth.\n\n"
            f"{style_instruction}\n\n"
            f"{context}"
            f"Question: {user_input.strip()}\n\nAnswer:"
        )

        start_time = time.time()
        with torch.no_grad():
            generated = pipe(
                generated_prompt,
                max_new_tokens=300,
                do_sample=True,
                top_p=0.9,
                temperature=0.7
            )
        elapsed_time = time.time() - start_time

        output_text = generated[0]["generated_text"]
        match = re.search(r"Answer:\s*(.*)", output_text, re.DOTALL)
        model_answer = match.group(1).strip() if match else output_text.strip()

        for kw in ["your students", "may also want", "in order to"]:
            if kw in model_answer:
                model_answer = model_answer.split(kw)[0].strip()
                break

        context_memory["last_question"] = user_input
        context_memory["last_answer"] = model_answer

        minutes = int(elapsed_time // 60)
        seconds = round(elapsed_time % 60, 1)
        time_display = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"

        return f"[Generated in {time_display}]\n\n{model_answer}"

    except Exception as e:
        return f"An error occurred:\n\n{e}"

# Gradio UI
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=5, placeholder="Ask your question here...", label="Ask your question"),
    outputs=gr.Textbox(lines=1, label="Response"),
    title="Gemma Mentor Chatbot",
    description="Ask questions and get warm, supportive, emotionally intelligent responses. Type 'exit' to reset.",
    flagging_mode="never",   # ✅ updated from deprecated allow_flagging
    theme="default",
    live=False,
    api_name="/predict"
)

# Correct launch block for Gradio 4+ / 5+
if __name__ == "__main__":
    interface.queue()             # ✅ enable queuing system
    interface.launch(share=False)