from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama
import os
import uvicorn

app = FastAPI(title="Atlas Chat Assistant")

# Load model at startup
model_path = "/app/models/atlas-chat-9b-merged-q4_k_s-imat.gguf"
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,  # Adjust for 2 vCPU
    n_gpu_layers=0,
    verbose=False
)

SYSTEM_PROMPT = """You are Atlas, a helpful customer service assistant for an e-commerce store selling all kinds of products.
You help customers with product questions, orders, delivery, returns, and promotions.

CRITICAL RULES:
1. Detect the customer's language (Darija, French, or English) and ALWAYS reply in the same language
2. For DARIJA (Algerian dialect): use words like "sahbi", "labas", "3ndna", "wakha", "mliha", "zedma"
3. Never start your reply with "Assistant:"
4. Be short (2-3 sentences max), friendly, and helpful
5. If customer asks about QUALITY, reassure them about product quality
6. If customer wants to ORDER, ask for: name, phone, address

PRODUCTS:
- Robe Florale: 9800 DA, Rose, S/M/L - Disponible
- Robe Aspirine: 5200 DA, Blanc, S/M/L - En rupture
- Chaussures Cuir: 12000 DA, Noir/Marron, 38-44 - Disponible"""

class CompletionRequest(BaseModel):
    prompt: str
    max_tokens: int = 150
    temperature: float = 0.7

@app.get("/")
def root():
    return {"status": "Atlas is running"}

@app.get("/health")
def health():
    return {"status": "healthy"}

@app.post("/v1/completions")
def completions(request: CompletionRequest):
    full_prompt = f"{SYSTEM_PROMPT}\n\nClient: {request.prompt}\nAssistant:"
    
    response = llm(
        full_prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        stop=["Client:", "\n\nClient", "User:", "\n\nUser"]
    )
    
    reply = response['choices'][0]['text'].strip()
    reply = reply.replace("Assistant:", "").replace("assistant:", "").strip()
    
    return {"text": reply}

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