"""Answer quality of a served config against reference answers from another config.

    python scripts/quality.py URL make REF.json    greedy reference answers (run on the exact config)
    python scripts/quality.py URL score REF.json   how likely this server finds those answers

Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
finished assistant message renders differently from the generation prompt, so re-scoring
through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
token (the reference config scores about its own decode logprobs); `same_tokens` is how many
leading tokens this server's own greedy answer shares with the reference.
"""

import argparse
import json
import urllib.request

PROMPTS = [
    "Write a Python function that returns the n-th Fibonacci number iteratively.",
    "Explain what `git rebase -i HEAD~3` does, in three sentences.",
    "Write a bash one-liner that counts the lines of all .py files below the current directory.",
    "On Gentoo, how do I rebuild every package that links against a library I just updated?",
    "What is the difference between a process and a thread? Answer in two sentences.",
]
CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."


def post(url, path, body):
    request = urllib.request.Request(url + path, json.dumps(body).encode(),
                                     {"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=3600) as response:
        return json.load(response)


def greedy(url, model, ids, tokens):
    choice = post(url, "/v1/completions", {
        "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
        "return_tokens_as_token_ids": True})["choices"][0]
    return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]


def nll(url, model, ids, answer):
    """Mean negative log-likelihood of the answer tokens after the prompt ids."""
    entries = post(url, "/v1/completions", {
        "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
        "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
    return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("url")
    parser.add_argument("action", choices=("make", "score"))
    parser.add_argument("ref")
    parser.add_argument("--model", default="gemma4")
    parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
    parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
    parser.add_argument("--code-chars", type=int, default=900)
    args = parser.parse_args()
    if args.action == "make":
        with open(args.code) as f:
            prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
        refs = []
        for prompt in prompts:
            ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
                                               "messages": [{"role": "user", "content": prompt}]})
            ids = ids["tokens"]
            text, answer = greedy(args.url, args.model, ids, args.tokens)
            refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
            print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
        with open(args.ref, "w") as f:
            json.dump(refs, f)
        return
    with open(args.ref) as f:
        refs = json.load(f)
    scores, same = [], []
    for ref in refs:
        scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
        text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
        same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
                         min(len(answer), len(ref["answer"]))))
        print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
                          "same_tokens": same[-1], "answer": text[:100]}), flush=True)
    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
                      "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))


if __name__ == "__main__":
    main()