from llama_cpp import Llama
import gradio as gr
import os
from datetime import datetime


# YOUR APP SETTINGS 
APP_NAME = "ChennaiAI" # Change this to your app name
MODEL_PATH = r"C:\Users\Elakeya\OneDrive\Desktop\Phi-3-mini-4k-instruct-q4.gguf"
LOGO = "🤖" # Change emoji

# LOAD MODEL 
print(f"Loading {APP_NAME}...")
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_threads=8,
    n_gpu_layers=0, # Set 35 for Nvidia GPU
    verbose=False
)

def ai_reply(message, history):
    history = history or []

    # System prompt = This makes your app behave differently
    system_prompt = f"<|system|>You are {APP_NAME}, a helpful AI assistant made in India.<|end|>\n"

    prompt = system_prompt
    for msg_obj in history:
        role = msg_obj.get("role")
        content = msg_obj.get("content", "")
        if role == "user":
            prompt += f"<|user|>{content}<|end|>\n"
        elif role == "assistant":
            prompt += f"<|assistant|>{content}<|end|>\n"

    prompt += f"<|user|>{message}<|end|>\n<|assistant|>"

    output = llm(prompt, max_tokens=512, temperature=0.7, stop=["<|user|>", "<|end|>"])
    response = output['choices'][0]['text'].strip()

    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})
    return "", history

def save_chat(history):
    if not history: 
        return "No chat to save"
    filename = f"chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
    with open(filename, "w", encoding="utf-8") as f:
        for msg_obj in history:
            role = msg_obj.get("role", "unknown").capitalize()
            content = msg_obj.get("content", "")
            f.write(f"{role}: {content}\n\n")
    return f"Saved as {filename}"

# YOUR APP UI (Gradio 6.0 compatible)
with gr.Blocks(title=APP_NAME, theme=gr.themes.Soft()) as app:
    gr.Markdown(f"# {LOGO} {APP_NAME}")
    gr.Markdown("Your own offline AI assistant. Built with Python.")

    chatbot = gr.Chatbot(height=500, label="Chat")
    
    with gr.Row():
        msg = gr.Textbox(label="Ask me anything", placeholder="Type here...", scale=4)
        send = gr.Button("Send", scale=1)

    with gr.Row():
        clear = gr.Button("Clear Chat")
        save = gr.Button("Save Chat")

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

    send.click(ai_reply, [msg, chatbot], [msg, chatbot])
    msg.submit(ai_reply, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot)
    save.click(save_chat, chatbot, status)
    app.launch(share=True)