"""Run a Jebadiah GGUF through llama.cpp's llama-server and print typed answers.

The prompt is rendered in Python exactly as AINode's /v1/systemone does (jebadiah_prompt.py, the chat
template with thinking off) and the rendered text goes to llama-server's raw /completion endpoint, so the
server's own chat template never touches it. Nothing is generated: the answer is read off the log
probabilities of the single-token option labels ("A", "B", ...) at the answer position, renormalised over
those labels, with the model's per-type temperature from temperatures.json applied.

  llama-server -m jebadiah-27b-Q4_K_M.gguf -c 4096 -np 1 --port 8080
  python scripts/decide_gguf.py --server http://127.0.0.1:8080 --request scripts/example-request.json

--tokenizer is a folder (or Hub repo id) with tokenizer.json, tokenizer_config.json and chat_template.jinja;
this repository ships them, so the default is the folder above scripts/. Needs `transformers` (the tokenizer
only, no torch) and nothing else outside the standard library.
"""
from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from jebadiah_prompt import Renderer, answer_from_probs  # noqa: E402


def load_renderer(tokenizer: str, max_tokens: int = 2048) -> Renderer:
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(tokenizer)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    return Renderer(tok, max_tokens)


def read_temperatures(path: str | None) -> dict:
    if not path or not os.path.exists(path):
        return {}
    return {k: float(v) for k, v in json.load(open(path))["temperatures"].items()}


def post(server: str, path: str, body: dict, timeout: float = 900) -> dict:
    req = urllib.request.Request(server.rstrip("/") + path, data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def label_logprobs(server: str, prompt: str, cand_ids: list[int], n_probs: int = 1000) -> tuple[list[float], int]:
    """Log probabilities (full-vocab softmax, before any sampling) of each candidate token at the position
    right after `prompt`. A label outside the top n_probs gets the smallest returned value, an upper bound
    that is already negligible after renormalisation. Returns (logprobs, number of labels not returned)."""
    r = post(server, "/completion", {"prompt": prompt, "n_predict": 1, "n_probs": n_probs,
                                     "post_sampling_probs": False, "cache_prompt": False,
                                     "temperature": 0.0})
    top = r["completion_probabilities"][0]["top_logprobs"]
    lp = {t["id"]: t["logprob"] for t in top}
    floor = min(lp.values())
    return [lp.get(c, floor) for c in cand_ids], sum(1 for c in cand_ids if c not in lp)


def option_probs(logprobs: list[float], temperature: float = 1.0) -> list[float]:
    """Softmax over the labels only, after dividing by the temperature. log p = logit - logsumexp(all
    logits), and the constant cancels in the softmax, so this equals softmax(logits[labels] / T)."""
    z = [x / temperature for x in logprobs]
    m = max(z)
    e = [math.exp(x - m) for x in z]
    s = sum(e)
    return [x / s for x in e]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--server", default="http://127.0.0.1:8080", help="a running llama-server with a Jebadiah GGUF")
    ap.add_argument("--request", required=True, help="JSON file: {state, questions: {id: {type, instructions, criteria}}}")
    ap.add_argument("--tokenizer", default=os.path.dirname(HERE), help="folder or Hub repo id with the tokenizer and chat template")
    ap.add_argument("--temperatures", default=os.path.join(os.path.dirname(HERE), "temperatures.json"))
    ap.add_argument("--no-temperatures", action="store_true", help="raw probabilities, as the served route returns today")
    ap.add_argument("--n-probs", type=int, default=1000)
    a = ap.parse_args()
    req = json.load(open(a.request))
    renderer = load_renderer(a.tokenizer)
    temps = {} if a.no_temperatures else read_temperatures(a.temperatures)
    out = {"temperatures_applied": temps, "answers": {}}
    for qid, q in req["questions"].items():
        rd = renderer.render(req["state"], q)
        lps, _ = label_logprobs(a.server, rd.prompt, rd.cand_ids, a.n_probs)
        probs = option_probs(lps, float(temps.get(q["type"], 1.0)))
        out["answers"][qid] = answer_from_probs(q, rd.keys, probs)
    print(json.dumps(out, indent=1))


if __name__ == "__main__":
    main()