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

# =========================
# LOAD MODEL
# =========================
model_path = hf_hub_download(
    repo_id="Mikecode123/ALX",
    filename="qwen2-1_5b-instruct-q4_0.gguf",
    token=os.getenv("HF_TOKEN")
)

llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2
)

# =========================
# CHAT FUNCTION (FIXED)
# =========================
def chat(message, history):
    # convert Gradio tuple history -> messages format
    messages = []

    for user, bot in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})

    messages.append({"role": "user", "content": message})

    output = llm.create_chat_completion(
        messages=messages,
        max_tokens=300,
        temperature=0.7
    )

    reply = output["choices"][0]["message"]["content"]

    history.append((message, reply))

    return history, history

# =========================
# UI (FIXED GRADIO 4)
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Living Legend AI Chat")

    chatbot = gr.Chatbot()
    msg = gr.Textbox()

    state = gr.State([])

    def respond(message, history):
        new_history, updated_state = chat(message, history)
        return "", new_history, updated_state

    msg.submit(respond, [msg, state], [msg, chatbot, state])

demo.launch()