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

# Download GGUF model from Hugging Face
model_path = hf_hub_download(
    repo_id="ikramknd/atlas-chat-9b-merged-IQ3_XXS-GGUF",
    filename="atlas-chat-9b-merged-iq3_xxs-imat.gguf"
)

# Load model with low RAM settings
llm = Llama(
    model_path=model_path,
    n_ctx=256,
    n_gpu_layers=0,
    n_threads=2,
    n_batch=32,
    use_mmap=True,
    use_mlock=False,
    chat_format="chatml",
    verbose=False,
)

# System prompt
SYSTEM_PROMPT = """
Assistant e-commerce jza2iri.
Respond in SAME language as client (darja/French/English).
Short, friendly answers.
Help with orders, products, delivery, returns.
"""

# Chat function
def chat(message, history):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT}
    ]

    # Keep only recent history
    for h in history[-4:]:
        if h[0]:
            messages.append({"role": "user", "content": str(h[0])[:80]})
        if h[1]:
            messages.append({"role": "assistant", "content": str(h[1])[:80]})

    # Current user message
    messages.append({
        "role": "user",
        "content": str(message)[:80]
    })

    # Generate response
    output = llm.create_chat_completion(
        messages=messages,
        max_tokens=60,
        temperature=0.6,
        top_k=30,
        top_p=0.85,
        stop=["<|im_end|>", "User:", "\n\n"],
    )

    response = output["choices"][0]["message"]["content"].strip()

    return response

# Gradio UI
demo = gr.ChatInterface(
    fn=chat,
    title="Atlas Algerian E-Commerce Assistant 🛍️",
    description="Parle darja, français, ou anglais",
)

# Better stability on HF Spaces
demo.queue()

# Launch app
demo.launch(server_name="0.0.0.0", server_port=7860)