"""Frozen teacher-forced perplexity and fixed-output serving throughput."""
import argparse
import collections
import concurrent.futures
import json
import math
import statistics
import re
import time
import urllib.request
from common import RUN, records, read_json, write_json, sha256
from evaluate import stream, post, key


def perplexity(api):
    windows = records(RUN / "evaluation/perplexity.jsonl")
    rows = []
    for i, window in enumerate(windows):
        with post(api, "/completions", {"model": "qwen3.8-27b", "prompt": window["text"],
                  "max_tokens": 1, "temperature": 0, "prompt_logprobs": 0}) as response:
            result = json.load(response)
        values = []
        for entry in result["choices"][0]["prompt_logprobs"][1:]:
            if entry is None or len(entry) != 1:
                raise ValueError("Expected exactly the observed token's log probability")
            lp = next(iter(entry.values()))
            lp = lp["logprob"] if isinstance(lp, dict) else lp
            assert math.isfinite(lp)
            values.append(lp)
        assert values
        rows.append({"id": i, "language": window["language"], "tokens": len(values), "logprob_sum": sum(values)})
    groups = collections.defaultdict(list)
    for r in rows:
        groups[r["language"]].append(r)
        groups["all"].append(r)
    return {"input_sha256": sha256(RUN / "evaluation/perplexity.jsonl"), "windows": rows,
            "scores": {name: {"tokens": sum(r["tokens"] for r in group),
                                "ppl": math.exp(-sum(r["logprob_sum"] for r in group)/sum(r["tokens"] for r in group))}
                       for name, group in groups.items()}}


def metrics(api):
    req = urllib.request.Request(api.removesuffix("/v1") + "/metrics", headers={"Authorization": "Bearer " + key()})
    with urllib.request.urlopen(req, timeout=10) as response:
        return response.read().decode()


def throughput(api, concurrency):
    prompts = ["Explain how a hash table handles collisions, with examples.",
               "Write a Python implementation of merge sort and explain its complexity.",
               "Describe how to design a reliable background job queue.",
               "Explain photosynthesis and the role of chlorophyll."]
    def one(i):
        return stream(api, {"model": "qwen3.8-27b", "messages": [{"role": "user", "content": prompts[i % len(prompts)]}],
                            "temperature": 0, "seed": 15027, "max_tokens": 512, "ignore_eos": True,
                            "chat_template_kwargs": {"enable_thinking": False}})
    one(0)  # identical warm-up for every model
    before = metrics(api)
    start = time.monotonic()
    with concurrent.futures.ThreadPoolExecutor(concurrency) as pool:
        rows = list(pool.map(one, range(max(8, concurrency*2))))
    seconds = time.monotonic()-start
    after = metrics(api)
    def counters(text):
        totals = collections.Counter()
        for line in text.splitlines():
            match = re.match(r'(vllm:spec_decode_num_(?:drafts|draft_tokens|accepted_tokens)_total)(?:\{[^}]*\})?\s+([\d.eE+-]+)',line)
            if match:
                totals[match[1]] += float(match[2])
        return totals
    delta = counters(after)
    delta.subtract(counters(before))
    assert all(r["usage"]["completion_tokens"] == 512 for r in rows), "Fixed-length benchmark stopped early"
    return {"concurrency": concurrency, "requests": len(rows), "output_tokens_per_request": 512,
            "wall_seconds": seconds, "aggregate_output_tps": sum(r["usage"]["completion_tokens"] for r in rows)/seconds,
            "median_decode_tps": statistics.median(r["decode_tps"] for r in rows),
            "median_ttft_seconds": statistics.median(r["ttft_seconds"] for r in rows),
            "sampling": "greedy; ignore_eos for this throughput test only", "calls": rows,
            "speculation_counter_deltas": dict(delta),
            "decode_tps_note": "Client stream timing estimate; speculative decoding may deliver several tokens in one event.",
            "metrics_before": before, "metrics_after": after}


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("tag")
    ap.add_argument("--api", default="http://127.0.0.1:18021/v1")
    ap.add_argument("--kind", choices=["ppl", "speed"], required=True)
    ap.add_argument("--concurrency", type=int, default=1)
    args = ap.parse_args()
    result = perplexity(args.api) if args.kind == "ppl" else throughput(args.api, args.concurrency)
    target = RUN / "results" / args.tag / ("perplexity.json" if args.kind == "ppl" else f"speed-c{args.concurrency}.json")
    write_json(target, result)
    print(target, flush=True)