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

# =========================
# LOAD MODEL FROM HF REPO
# =========================
model_path = hf_hub_download(
    repo_id="Mikecode123/ALX",
    filename="qwen2-1_5b-instruct-q4_0.gguf"
)

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

# =========================
# CHAT FUNCTION
# =========================
def chat(prompt, history):
    messages = []

    # convert gradio history to chat format
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})

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

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

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

    history.append((prompt, response))
    return "", history

# =========================
# GRADIO UI
# =========================
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Qwen GGUF AI (Living Legend Build)")

    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Ask your AI")
    clear = gr.Button("Clear")

    msg.submit(chat, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot)

# =========================
# LAUNCH
# =========================
demo.launch()