"""Needle reflex sidecar — stdio JSONL server around Cactus Needle (26M
on-device tool-calling model, MIT). Spawned by src/koi/reflex/needle-reflex.ts
and kept warm: model + JIT stay loaded so repeated calls avoid startup cost.

Protocol (one JSON object per line):
  in : {"id": 1, "text": "turn it down", "tools": [...]}    (tools = JSON array)
  out: {"id": 1, "calls": [{"name": "...", "arguments": {...}}], "ms": 812}
  out: {"id": 1, "error": "..."}                            (never crashes the loop)
A {"id": 0, "warmup": true} request runs a dummy generate to pay JIT upfront.

Run with the needle checkout's venv python, cwd = the needle repo (so its
checkpoints/ and tokenizer/ resolve): see reflex config in skykoi.json.
"""
import json
import os
import sys
import time

# The script lives in the skykoi repo but runs with cwd = the needle checkout;
# python puts the SCRIPT's dir on sys.path, not cwd — add it explicitly.
sys.path.insert(0, os.getcwd())

sys.stderr.write("needle-sidecar: loading model...\n")
sys.stderr.flush()

from needle.model.run import load_checkpoint, generate  # noqa: E402
from needle.model.architecture import SimpleAttentionNetwork  # noqa: E402
from needle.dataset.dataset import get_tokenizer  # noqa: E402

_t0 = time.time()
PARAMS, CONFIG = load_checkpoint("checkpoints/checkpoint_epoch3.pkl")
MODEL = SimpleAttentionNetwork(CONFIG)
TOKENIZER = get_tokenizer()


def fast_generate(query, tools_json):
    return generate(
        MODEL, PARAMS, TOKENIZER, query, tools=tools_json, stream=False,
        max_enc_len=384, max_gen_len=128,
    )


def validated_calls(raw_text, tools):
    """Fail-closed replacement for per-step constrained decoding: only calls
    naming a KNOWN tool with plausible JSON arguments survive."""
    try:
        calls = json.loads(raw_text)
    except Exception:  # noqa: BLE001
        return []
    if not isinstance(calls, list):
        return []
    known = {t.get("name"): t.get("parameters", {}) for t in tools if isinstance(t, dict)}
    ok = []
    for c in calls:
        if not isinstance(c, dict) or c.get("name") not in known:
            continue
        args = c.get("arguments")
        if not isinstance(args, dict):
            continue
        params = known[c.get("name")] or {}
        if any(k not in params for k in args):
            continue
        required = [k for k, v in params.items() if isinstance(v, dict) and v.get("required")]
        if any(k not in args for k in required):
            continue
        ok.append({"name": c["name"], "arguments": args})
    return ok


sys.stderr.write(f"needle-sidecar: ready in {time.time() - _t0:.1f}s\n")
sys.stderr.flush()
print(json.dumps({"ready": True}), flush=True)

for line in sys.stdin:
    line = line.strip()
    if not line:
        continue
    req_id = None
    try:
        req = json.loads(line)
        req_id = req.get("id")
        text = str(req.get("text", "")) if not req.get("warmup") else "hello there"
        tools = req.get("tools", [])
        if isinstance(tools, str):
            tools = json.loads(tools or "[]")
        t1 = time.time()
        raw = fast_generate(text, json.dumps(tools))
        calls = validated_calls(raw, tools)
        print(
            json.dumps({"id": req_id, "calls": calls, "ms": round((time.time() - t1) * 1000)}),
            flush=True,
        )
    except Exception as err:  # noqa: BLE001 — the loop must survive anything
        print(json.dumps({"id": req_id, "error": str(err)[:300]}), flush=True)