"""Modal deployment: llama.cpp serving Qwen3.5-9B with JSON schema
enforcement (ARCHITECTURE.md D3). Built on the official llama.cpp CUDA image.

Deploy:  modal deploy modal_app/inference.py
Secret:  modal secret create bds-auth BDS_TOKEN=<random hex>
Then set on the HF Space / locally:
  MODAL_URL=<the printed generate endpoint url>
  MODAL_TOKEN=<same hex>
"""
from __future__ import annotations

import json
import os
import subprocess
import time
import urllib.request

import modal

# Qwen3.5-9B Q4_K_M ≈ 6GB — newer generation, better multi-turn coherence
# for board presentations, inside the 16GB budget. (Forfeits Tiny Titan ≤4B.)
MODEL_REPO = "bartowski/Qwen_Qwen3.5-9B-GGUF"
MODEL_FILE = "Qwen_Qwen3.5-9B-Q4_K_M.gguf"
LLAMA_PORT = 8081

image = (
    modal.Image.from_registry(
        "ghcr.io/ggml-org/llama.cpp:server-cuda", add_python="3.11")
    .entrypoint([])  # the image defaults to exec llama-server; we manage it
    .pip_install("fastapi[standard]", "huggingface_hub")
    .run_commands(
        "python -c \"from huggingface_hub import hf_hub_download; "
        f"hf_hub_download('{MODEL_REPO}', '{MODEL_FILE}', "
        "local_dir='/models')\""
    )
)

app = modal.App("brad-did-something", image=image)

with image.imports():
    from fastapi import HTTPException, Request


def _find_server_binary() -> str:
    for path in ("/app/llama-server", "/llama-server",
                 "/usr/local/bin/llama-server"):
        if os.path.exists(path):
            return path
    return "llama-server"  # hope it's on PATH


@app.cls(
    gpu="L4",
    scaledown_window=300,  # stay warm between calls within a play session
    timeout=120,
    secrets=[modal.Secret.from_name("bds-auth")],
)
class Llama:
    @modal.enter()
    def start_server(self):
        self.proc = subprocess.Popen([
            _find_server_binary(),
            "--model", f"/models/{MODEL_FILE}",
            "--ctx-size", "4096",
            "--n-gpu-layers", "99",
            "--port", str(LLAMA_PORT),
            "--host", "127.0.0.1",
        ])
        deadline = time.time() + 120
        while time.time() < deadline:
            try:
                urllib.request.urlopen(
                    f"http://127.0.0.1:{LLAMA_PORT}/health", timeout=2)
                return
            except Exception:
                time.sleep(1)
        raise RuntimeError("llama-server did not become healthy")

    @modal.exit()
    def stop_server(self):
        self.proc.terminate()

    @modal.fastapi_endpoint(method="POST")
    def generate(self, body: dict, request: Request):
        expected = os.environ.get("BDS_TOKEN", "")
        sent = request.headers.get("authorization", "")
        if expected and sent != f"Bearer {expected}":
            raise HTTPException(401, "bad token")

        t0 = time.time()
        # the empty <think> block disables Qwen3.5's default thinking mode —
        # the JSON grammar takes over immediately after
        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"
        )
        payload = json.dumps({
            "prompt": prompt,
            "temperature": 0.4,
            # 512 truncated the JSON once crises/events gained the long
            # image_prompt + comic_caption fields → unterminated-string parse
            # failures → fallbacks. 1024 leaves comfortable headroom.
            "n_predict": 1024,
            "cache_prompt": True,
            "json_schema": body["schema"],  # grammar-enforced at generation
        }).encode()
        req = urllib.request.Request(
            f"http://127.0.0.1:{LLAMA_PORT}/completion",
            data=payload, headers={"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=90) as resp:
                out = json.loads(resp.read())
            data = json.loads(out["content"])
            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)[:200],
                    "ms": int((time.time() - t0) * 1000)}