"""Local real-AI inference server — same /generate contract as the Modal app,
same runtime (llama.cpp) and same Qwen3.5-9B GGUF, just on CPU for development.

Usage:
  python tests/local_llm_server.py [--port 8081] [--model models/<file>.gguf]
Then point the game at it:
  set MODAL_URL=http://127.0.0.1:8081
  set BDS_LLM_TIMEOUT=90
  python app.py
"""
from __future__ import annotations

import argparse
import json
import time

import uvicorn
from fastapi import FastAPI

parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8081)
parser.add_argument("--model", default="models/Qwen_Qwen3.5-9B-Q4_K_M.gguf")
args = parser.parse_args()

print(f"loading {args.model} ...")
from llama_cpp import Llama, LlamaGrammar  # noqa: E402  (slow import)
from llama_cpp.llama_grammar import json_schema_to_gbnf  # noqa: E402

llm = Llama(model_path=args.model, n_ctx=4096, n_threads=None, verbose=False)
print("model loaded.")

app = FastAPI()


@app.post("/generate")
def generate(body: dict):
    t0 = time.time()
    try:
        schema = body["schema"]
        grammar = LlamaGrammar.from_string(
            json_schema_to_gbnf(json.dumps(schema)), verbose=False)
        # mirror the Modal path: the empty <think> block disables Qwen3.5's
        # thinking mode so the JSON grammar takes over immediately
        prompt = (
            f"<|im_start|>system\n{body['system_prompt']}\n"
            f"GAME STATE JSON:\n{json.dumps(body.get('context', {}))}\n<|im_end|>\n"
            f"<|im_start|>user\n{body['user_prompt']}<|im_end|>\n"
            f"<|im_start|>assistant\n<think>\n\n</think>\n\n"
        )
        out = llm(prompt, max_tokens=1024, temperature=0.4, grammar=grammar)
        text = out["choices"][0]["text"]
        data = json.loads(text)
        return {"ok": True, "data": data, "ms": int((time.time() - t0) * 1000)}
    except Exception as exc:  # caller falls back; never crash the endpoint
        return {"ok": False, "error": str(exc)[:300],
                "ms": int((time.time() - t0) * 1000)}


@app.get("/healthz")
def healthz():
    return {"ok": True}


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=args.port)