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

APP_NAME = "ChennaiAI"
LOGO = "🤖"

MODEL_PATH = hf_hub_download(
    repo_id="microsoft/Phi-3-mini-4k-instruct-gguf",
    filename="Phi-3-mini-4k-instruct-q4.gguf"
)

# LOAD MODEL
print(f"Loading {APP_NAME}...")
llm = Llama(
    model_path=MODEL_PATH,
    n_ctx=4096,
    n_threads=2,  # Set to 2 vCPUs for Hugging Face free tier
    n_gpu_layers=0,
    verbose=False
)

def ai_reply(message, history):
    history = history or []
    
    # Format prompt from list of tuples [(user, assistant)]
    prompt = f"<|system|>You are {APP_NAME}, a helpful AI assistant.<|end|>\n"
    for user_msg, bot_msg in history:
        prompt += f"<|user|> {user_msg} <|end|>\n<|assistant|> {bot_msg} <|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()
    
    # Append as tuple format required by Gradio
    history.append((message, 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 user_msg, bot_msg in history:
            f.write(f"User: {user_msg}\nAssistant: {bot_msg}\n\n")
    return f"Saved as {filename}"

# YOUR APP UI
with gr.Blocks(title=APP_NAME, theme=gr.themes.Soft()) as app:
    gr.Markdown(f"# {LOGO} {APP_NAME}")
    gr.Markdown("Your own 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()