import os, json
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama

MODEL_PATH = os.path.join(os.environ.get("MODEL_DIR", "/home/user/models"),
                          os.environ.get("MODEL_FILE", "astra-meal-parser-q4_k_m.gguf"))
API_KEY = os.environ.get("API_KEY", "")   # set as a Space secret

SYSTEM = (
    "You are a meal parser. Extract every food item and its amount from the user's "
    "meal description (Turkish or English). Return ONLY a strict JSON object of the form "
    '{"items": [{"name": string, "amount": string}]}. '
    "No macros, no calories, no conversational text, no markdown, only valid JSON."
)

llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=2, chat_format="chatml")
app = FastAPI()

class MealIn(BaseModel):
    meal: str

def extract_json(text):
    s, e = text.find("{"), text.rfind("}")
    if s == -1 or e == -1:
        return None
    try:
        return json.loads(text[s:e+1])
    except Exception:
        return None

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

@app.post("/parse")
def parse(body: MealIn, authorization: str = Header(default="")):
    if API_KEY and authorization != f"Bearer {API_KEY}":
        raise HTTPException(status_code=401, detail="Unauthorized")
    out = llm.create_chat_completion(
        messages=[{"role": "system", "content": SYSTEM},
                  {"role": "user", "content": body.meal}],
        temperature=0, max_tokens=256, stop=["<|im_end|>"],
    )
    parsed = extract_json(out["choices"][0]["message"]["content"])
    if parsed is None:
        raise HTTPException(status_code=422, detail="Parse failed")
    return parsed