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

# Load GPT-Neo 1.3B
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")

# Chat history as list of tuples (user, ai)
chat_history = []

def chat(message, history):
    # Append user message
    history.append((message, ""))  # placeholder for AI response

    # Take last 3 user messages for context
    recent_user_msgs = [m[0] for m in history[-3:]]
    context = "\n".join(recent_user_msgs) + "\nAI: "

    inputs = tokenizer(context, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        do_sample=True,
        top_p=0.9,
        temperature=0.8
    )

    # Decode and join sentences to avoid one-word-per-line
    raw_response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    sentences = re.split(r'(?<=[.!?]) +', raw_response)
    response = ' '.join(sentences)

    # Update last tuple with AI response
    history[-1] = (message, response)

    # Return updated chat history
    return history

# Gradio interface with spooky theme
with gr.Blocks(css="""
body {
    background-color: #0a1f44; /* Deep blue background */
}
h1 {
    color: red;
    font-family: 'Creepster', cursive;
    text-align: center;
    font-size: 48px;
    margin-bottom: 20px;
}
.chatbot .message.ai {
    background-color: #4444ff;
    color: white;
    border-radius: 10px;
    padding: 8px;
}
.chatbot .message.user {
    background-color: #2222aa;
    color: white;
    border-radius: 10px;
    padding: 8px;
}
""") as demo:

    gr.HTML("<h1>💀 Welcome to Ghost AI 💀</h1>")
    chatbot = gr.Chatbot()
    txt = gr.Textbox(placeholder="Type your message here...")
    
    # Submit user input to chat function
    txt.submit(chat, [txt, chatbot], chatbot)

demo.launch()