import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Pick a small open-source model for demo (can change later, e.g., "mistralai/Mistral-7B-Instruct-v0.2")
MODEL_NAME = "microsoft/DialoGPT-medium"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Chat history (stored in memory)
chat_history = []

# System persona
SYSTEM_PROMPT = ""

def chat_with_gourav(user_input):
    global chat_history

    # Append system + user input
    chat_history.append(f"User: {user_input}")
    prompt = SYSTEM_PROMPT + "\n" + "\n".join(chat_history) + "\nSmartEmail Assisstant:"

    # Tokenize
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=500,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        temperature=0.7,
        top_p=0.9
    )

    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    chat_history.append(f"Gourav: {response}")

    return response, chat_history


# Reset button
def reset_chat():
    global chat_history
    chat_history = []
    return "Chat history cleared."


# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 SmartEmail Assisstant")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type your message...")
    clear = gr.Button("Clear Chat")

    def user(user_message, history):
        response, updated_history = chat_with_gourav(user_message)
        messages = []
        for h in updated_history:
            if h.startswith("User:"):
                messages.append([h.replace("User:", "").strip(), None])
            elif h.startswith("SmartEmail Assisstant:"):
                if len(messages) > 0:
                    messages[-1][1] = h.replace("SmartEmail Assisstant:", "").strip()
        return "", messages

    msg.submit(user, [msg, chatbot], [msg, chatbot])
    clear.click(reset_chat, None, chatbot)

demo.launch()