from fastapi import FastAPI, Request, HTTPException
from llama_cpp import Llama
import json
from json.decoder import JSONDecodeError

# (Your existing model path here) Please note this was my path to the downloaded model replace with your for it to work.
MODEL_PATH = r"C:\Users\YOUR_COMPUTER_NAME\.cache\huggingface\hub\models--TheBloke--stablelm-zephyr-3b-GGUF\snapshots\465310e5eb914190e89f531cc1c348812e27dcea\stablelm-zephyr-3b.Q5_K_M.gguf"
llm = Llama(model_path=MODEL_PATH, n_threads=4, n_ctx=512)

app = FastAPI()

@app.post("/insights")
async def insights(req: Request):
    raw = await req.body()                # get the raw bytes
    text_raw = raw.decode(errors="ignore")
    # Attempt a clean json.loads first
    try:
        payload = json.loads(text_raw)
    except JSONDecodeError:
        # Fallback: extract the first {...} block
        start = text_raw.find("{")
        end   = text_raw.rfind("}") + 1
        if start < 0 or end < 0 or end <= start:
            raise HTTPException(status_code=400, detail="Invalid JSON payload")
        try:
            payload = json.loads(text_raw[start:end])
        except JSONDecodeError:
            raise HTTPException(status_code=400, detail="Could not parse JSON payload")
    # At this point, payload is a dict
    prompt = payload.get("prompt", "")
    if not isinstance(prompt, str) or prompt.strip() == "":
        raise HTTPException(status_code=400, detail="Missing or empty 'prompt' field")
    # Generate AI insight
    out = llm(prompt, max_tokens=64)
    # out is a dict, extract text
    try:
        insight = out["choices"][0]["text"].strip()
    except Exception:
        raise HTTPException(status_code=500, detail="Model did not return valid choices")
    return {"insight": insight}