import gradio as gr
import torch
import json
import os
import time
from transformers import AutoTokenizer, AutoModelForCausalLM

# === CONFIGURATION ===
MODEL_NAME = "facebook/opt-350m"  # ✅ Fast and lightweight
MEMORY_FILE = "memory.json"

# === LOAD MODEL ===
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# === MEMORY ===
def load_memory(file_path):
    if os.path.exists(file_path):
        with open(file_path, "r", encoding="utf-8") as f:
            return json.load(f)
    return []

def save_memory(file_path, data):
    with open(file_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=4)

memory = load_memory(MEMORY_FILE)

# === CHAT FUNCTION ===
def chat_with_ai(user_input, history, turbo_mode):
    memory.append({"user": user_input})
    save_memory(MEMORY_FILE, memory)

    conversation = ""
    for msg in memory[-3:]:  # Only the last 3 inputs to keep it lightweight
        conversation += f"User: {msg['user']}\nAssistant:"

    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=512).to(model.device)

    generate_args = {
        'input_ids': inputs.input_ids,
        'max_length': inputs.input_ids.shape[-1] + (30 if turbo_mode else 50),
        'temperature': 0.6,
        'top_k': 30,
        'top_p': 0.9,
        'do_sample': True,
        'pad_token_id': tokenizer.eos_token_id,
        'early_stopping': True
    }

    output = model.generate(**generate_args)
    reply_tokens = output[0][inputs.input_ids.shape[-1]:]
    reply_text = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()

    # Save Assistant reply to memory
    memory.append({"assistant": reply_text})
    save_memory(MEMORY_FILE, memory)

    return "", history + [{"role": "user", "content": user_input}, {"role": "assistant", "content": reply_text}], reply_text

def clear_chat():
    memory.clear()
    save_memory(MEMORY_FILE, [])
    return "", [], ""

def save_chat(history):
    timestamp = time.strftime("%Y-%m-%d_%H-%M-%S")
    filename = f"chat_{timestamp}.txt"
    with open(filename, "w", encoding="utf-8") as f:
        for msg in history:
            f.write(f"{msg['role'].capitalize()}: {msg['content']}\n\n")
    return f"💾 Chat saved as {filename}."

# === DARK THEME CSS ===
css = """
body {
    background-color: #121212;
    color: #e0e0e0;
}
.gradio-container, .gr-block, .gr-button, .gr-textbox {
    background-color: #1e1e1e !important;
    border-color: #333 !important;
}
.gr-button:hover {
    background-color: #333333 !important;
}
"""

# === GRADIO APP ===
with gr.Blocks(css=css) as app:
    gr.Markdown("## 🤖 Ultra Fast Living AI - Clean Mode", elem_id="title")

    chatbot = gr.Chatbot(label="Chat History", type="messages", render=False)
    user_input = gr.Textbox(lines=3, label="Your Message", placeholder="Type your message...")
    ai_output = gr.Textbox(lines=5, label="🤖 AI's Response", interactive=False, placeholder="AI will respond here...")

    with gr.Row():
        send_btn = gr.Button("Send")
        clear_btn = gr.Button("Clear Chat")
        save_btn = gr.Button("Save Chat")
        turbo_mode = gr.Checkbox(label="🚀 Turbo Mode (Faster Replies)", value=True)

    status_box = gr.Textbox(label="System Status", interactive=False)

    send_btn.click(chat_with_ai, inputs=[user_input, chatbot, turbo_mode], outputs=[user_input, chatbot, ai_output])
    user_input.submit(chat_with_ai, inputs=[user_input, chatbot, turbo_mode], outputs=[user_input, chatbot, ai_output])
    clear_btn.click(clear_chat, outputs=[user_input, chatbot, ai_output])
    save_btn.click(save_chat, inputs=[chatbot], outputs=[status_box])

app.launch()