"""Correctness suite against one running server; dumps JSON for later A/B.

Usage: python3 validate.py <port> <model-name> <out.json>

Collects:
  1. 17*23 deterministic arithmetic check (temp 0).
  2. 700+ token generation; flags degenerate repetition (4-gram loop scan).
  3. 10 fixed greedy generations (50 tok, logprobs=5) - token ids + logprobs.
  4. prompt_logprobs teacher-forced scoring of 5 fixed paragraphs -
     per-token logprob under identical context (the clean A/B signal).
"""

import json
import sys
import urllib.request

PORT, MODEL, OUT = int(sys.argv[1]), sys.argv[2], sys.argv[3]
BASE = f"http://127.0.0.1:{PORT}"

GEN_PROMPTS = [
    "What is 17*23? Answer with just the number.",
    "Write a Python function that returns the nth Fibonacci number iteratively.",
    "Explain the difference between TCP and UDP in two sentences.",
    "Translate to French: 'The quick brown fox jumps over the lazy dog.'",
    "List the first 8 prime numbers separated by commas.",
    "What year did the Apollo 11 mission land on the moon? One word answer.",
    "Summarize the plot of Romeo and Juliet in one sentence.",
    "Write a SQL query selecting the top 5 customers by total order value from tables customers(id,name) and orders(id,customer_id,value).",
    "What is the derivative of x^3 + 2x with respect to x?",
    "Name the chemical symbol for gold and the element with atomic number 6.",
]

SCORE_TEXTS = [
    "The mitochondrion is the powerhouse of the cell, converting nutrients into adenosine triphosphate through oxidative phosphorylation. This process occurs across the inner mitochondrial membrane, where the electron transport chain establishes a proton gradient.",
    "def quicksort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quicksort(left) + middle + quicksort(right)",
    "In 1969, the Apollo 11 mission successfully landed the first humans on the Moon. Neil Armstrong and Buzz Aldrin spent approximately two and a quarter hours outside the spacecraft, collecting lunar material to bring back to Earth.",
    "Le petit prince demanda au renard ce que signifiait le mot apprivoiser. Le renard expliqua que cela signifiait creer des liens, et que si le prince l'apprivoisait, ils auraient besoin l'un de l'autre.",
    "The gradient of the loss function with respect to the weights is computed via backpropagation, applying the chain rule layer by layer from the output back to the input. Stochastic gradient descent then updates each weight proportionally.",
]


def post(path, payload, timeout=600):
    req = urllib.request.Request(
        BASE + path,
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def degenerate(text, n=4, thresh=12):
    words = text.split()
    if len(words) < n * thresh:
        return False
    grams = {}
    for i in range(len(words) - n):
        g = tuple(words[i : i + n])
        grams[g] = grams.get(g, 0) + 1
    return max(grams.values()) >= thresh


out = {"port": PORT, "model": MODEL}

# 1. arithmetic
r = post(
    "/v1/chat/completions",
    {
        "model": MODEL,
        "messages": [{"role": "user", "content": GEN_PROMPTS[0]}],
        "temperature": 0,
        "max_tokens": 2048,
    },
)
msg = r["choices"][0]["message"]
arith = (msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")
out["arith_correct"] = "391" in arith
out["arith_raw"] = (msg.get("content") or "")[:200]

# 2. long generation
r = post(
    "/v1/chat/completions",
    {
        "model": MODEL,
        "messages": [
            {
                "role": "user",
                "content": "Write a detailed 800-word essay on the history of container shipping, covering its origins, standardization, and global economic impact.",
            }
        ],
        "temperature": 0,
        "max_tokens": 2400,
    },
)
msg = r["choices"][0]["message"]
essay = (msg.get("reasoning_content") or "") + (msg.get("content") or "")
out["long_gen_tokens"] = r["usage"]["completion_tokens"]
out["long_gen_degenerate"] = degenerate(essay)
out["long_gen_tail"] = essay[-300:]

# 3. greedy generations with logprobs (completions API for raw control)
gens = []
for p in GEN_PROMPTS:
    r = post(
        "/v1/completions",
        {
            "model": MODEL,
            "prompt": p,
            "temperature": 0,
            "max_tokens": 50,
            "logprobs": 5,
        },
    )
    ch = r["choices"][0]
    gens.append(
        {
            "prompt": p,
            "tokens": ch["logprobs"]["tokens"],
            "token_logprobs": ch["logprobs"]["token_logprobs"],
        }
    )
out["greedy"] = gens

# 4. teacher-forced prompt logprobs
scored = []
for t in SCORE_TEXTS:
    r = post(
        "/v1/completions",
        {
            "model": MODEL,
            "prompt": t,
            "temperature": 0,
            "max_tokens": 1,
            "prompt_logprobs": 1,
        },
    )
    ch = r["choices"][0]
    plp = ch.get("prompt_logprobs")
    toks = []
    if plp:
        for pos in plp:
            if pos is None:
                toks.append(None)
                continue
            # dict: token_id -> {logprob, rank, decoded_token}; the entry with
            # rank field present; chosen token is the key matching prompt token
            entry = {
                tid: {"logprob": v["logprob"], "rank": v.get("rank")}
                for tid, v in pos.items()
            }
            toks.append(entry)
    scored.append({"text": t[:60], "prompt_logprobs": toks})
out["scored"] = scored

json.dump(out, open(OUT, "w"))
print(json.dumps({k: v for k, v in out.items() if k not in ("greedy", "scored")}, indent=1))
print(f"greedy prompts: {len(gens)}, scored texts: {len(scored)} -> {OUT}")