import gradio as gr
import os
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn

HF_TOKEN = os.environ.get("HF_TOKEN")

SYSTEM_PROMPT = """Tu es l'assistant SAV officiel du Centre Chery Tunisie.
Tu réponds en français ou en arabe dialectal tunisien selon la langue du client.
Réponds uniquement aux questions liées aux véhicules Chery."""

print("⏳ Downloading GGUF model...")
model_path = hf_hub_download(
    repo_id="dali4444444/chery-sav-chatbot-gguf",
    filename="chery-sav-chatbot-q4_k_m.gguf",
    token=HF_TOKEN
)

print("⏳ Loading model on CPU...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,
    verbose=False
)
print("✅ Model ready!")

def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for h in history:
        if isinstance(h, dict):
            messages.append({"role": h["role"], "content": h["content"]})
        else:
            messages.append({"role": "user", "content": h[0]})
            messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=300,
        temperature=0.7,
        top_p=0.9,
    )
    return response["choices"][0]["message"]["content"]

app = FastAPI()

class ChatRequest(BaseModel):
    message: str
    history: list = []

@app.post("/api/chat")
async def api_chat(req: ChatRequest):
    return {"reply": chat(req.message, req.history)}

demo = gr.ChatInterface(fn=chat, title="🚗 Chery SAV Assistant")
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)