import gradio as gr
from transformers import pipeline
import torch

# 1. SETUP
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
print("Loading Mochi... 0 Error Version.")

pipe = pipeline(
    "text-generation", 
    model=model_id, 
    device="cpu", 
    dtype=torch.float32 # Fixed the torch_dtype warning
)

# Persistent dictionary for separate users
user_memories = {}

def chat_logic(user_id, message):
    uid = str(user_id)
    if uid not in user_memories:
        user_memories[uid] = []
    
    system_prompt = "You are Mochi, a chill best friend. Talk like a 20-year-old. Keep it short."

    # Build history for the AI
    messages = [{"role": "system", "content": system_prompt}]
    
    # Only send last 8 messages to keep CPU fast
    for msg in user_memories[uid][-8:]:
        messages.append(msg)
    
    messages.append({"role": "user", "content": message})
    
    try:
        out = pipe(messages, max_new_tokens=100, do_sample=True, temperature=0.8, truncation=True)
        response = out[0]['generated_text'][-1]['content']
        
        # Save to memory
        user_memories[uid].append({"role": "user", "content": message})
        user_memories[uid].append({"role": "assistant", "content": response})
        
        # Keep internal memory small
        if len(user_memories[uid]) > 20:
            user_memories[uid] = user_memories[uid][-20:]
            
        return response
    except Exception as e:
        print(f"Error: {e}")
        return "brain lag lol. try again?"

# 2. THE WEB UI (Simplified for 0 errors)
with gr.Blocks() as demo:
    gr.Markdown("# 🐾 Mochi AI")
    
    user_id_input = gr.Textbox(label="User ID", value="default_user")
    # REMOVED type="messages" to stop the crash
    chatbot = gr.Chatbot(label="Chat with Mochi") 
    msg = gr.Textbox(label="Your Message", placeholder="Type here...")
    clear = gr.Button("Clear Chat")

    def respond(uid, message, chat_history):
        bot_message = chat_logic(uid, message)
        # For old Gradio, we append a [user, bot] list
        chat_history.append([message, bot_message])
        return "", chat_history

    msg.submit(respond, [user_id_input, msg, chatbot], [msg, chatbot])
    
    def clear_mem(uid):
        if uid in user_memories: user_memories[uid] = []
        return []
    
    clear.click(clear_mem, [user_id_input], [chatbot])

# 3. LAUNCH (Theme moved here to fix warning)
demo.launch(theme=gr.themes.Soft(primary_hue="pink"))