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

# Load model & tokenizer
model_id = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
model.eval()

# Memory for chat history
chat_history = []

# System prompt
system_prompt = "You are a helpful and friendly assistant."

def chat(user_input):
    # Combine system prompt and chat history
    prompt = system_prompt + "\n"
    for speaker, msg in chat_history:
        prompt += f"{speaker}: {msg}\n"
    prompt += f"User: {user_input}\nBot:"

    # Tokenize input with attention mask
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs.input_ids
    attention_mask = inputs.attention_mask

    with torch.no_grad():
        output = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=100,
            do_sample=True,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )

    full_reply = tokenizer.decode(output[0], skip_special_tokens=True)
    reply = full_reply.split("Bot:")[-1].strip()

    chat_history.append(("User", user_input))
    chat_history.append(("Bot", reply))

    return format_chat(chat_history)

def format_chat(history):
    return "\n\n".join([f"**{speaker}:** {msg}" for speaker, msg in history])

def clear():
    chat_history.clear()
    return ""

def export_chat():
    export = "\n".join([f"{s}: {m}" for s, m in chat_history])
    return gr.File.update(value=export.encode("utf-8"), filename="chat.txt")

with gr.Blocks() as demo:
    gr.Markdown("""# 🤖 Chat with Phi-2
Small CPU-friendly chatbot using Hugging Face Transformers
""")

    chatbot = gr.Markdown()
    with gr.Row():
        user_input = gr.Textbox(placeholder="Type your message here...", show_label=False)
        send_btn = gr.Button("Send")
    with gr.Row():
        clear_btn = gr.Button("🧹 Clear")
        export_btn = gr.Button("💾 Export Chat")
        file_out = gr.File(label="Download Chat")

    send_btn.click(chat, inputs=user_input, outputs=chatbot)
    user_input.submit(chat, inputs=user_input, outputs=chatbot)
    clear_btn.click(clear, outputs=chatbot)
    export_btn.click(export_chat, outputs=file_out)

    # Initialize chatbot display
    chatbot.value = format_chat(chat_history)

demo.launch(share=True)