#!/usr/bin/env python3
"""bench-vllm-corpus.py — build a code-shaped benchmark corpus from the repo.

Why this exists: `vllm bench serve --dataset-name random` produces synthetic
token soup, and speculative decoding accepts it nothing like real code — measured
on frank-medium, random prompts accepted 2.9-3.6 tokens/update against the ~5.8
the upstream reference reports on a code mix, which makes the random corpus
under-report throughput by a wide margin and would pick the wrong SPEC depth.

The prompts are the shape agents actually send: a repository artifact plus an
instruction that requires reading it. Output length is fixed at bench time
(`--custom-output-len`), so the corpus carries prompts only.

Usage:
    bench-vllm-corpus.py [repo-root] [out.jsonl]

Then feed it to the endpoint (see scripts/bench-vllm-serve.sh, which passes the
tokenizer the client needs and salts prompts so no run inherits a warm prefix
cache). Methodology and the numbers these corpora produced are in
serverAppConfigs-docs/LLM/frank-medium-radiance-mxfp4.md.
"""

from __future__ import annotations

import json
import pathlib
import sys

REPO = pathlib.Path(sys.argv[1] if len(sys.argv) > 1 else "/home/dank/serverAppConfigs")
OUT = pathlib.Path(sys.argv[2] if len(sys.argv) > 2 else "/home/dank/bench/corpus-code.jsonl")

TARGETS = [
    "scripts/deploy-remote.sh",
    "scripts/deploy-plan.py",
    "scripts/deploy.sh",
    "dankserver/frank-medium/radiance-entrypoint.sh",
    "dankserver/frank-medium/quadlet/frank-medium-a.container",
    "dankserver/frank-medium/proxy/affinity-proxy.py",
    "dankserver/frank-medium/proxy/routing.py",
    "tests/frank_medium_quadlet.bats",
    "tests/deployment_scoping.bats",
    "dankclaw/hermes/quadlet/hermes.container",
    "dankclaw/anubis/quadlet/anubis.container",
    "shared/prometheus/prometheus.yml",
    "shared/prometheus/rules/llm-fleet.yml",
    "shared/loki/loki-config.yaml",
    "dankclaw/compose.yaml",
    "dankserver/copyparty/quadlet/copyparty.container",
]

PROMPTS = [
    "Read this repository artifact and summarise what it does, then list every "
    "assumption it makes about the host it runs on.\n\n```\n{body}\n```",
    "Here is a file from our homelab repository. Identify any bug, race or "
    "missing error path, and show the minimal patch that fixes it.\n\n```\n{body}\n```",
    "Explain this configuration to a new operator, then rewrite it with "
    "comments that state the constraint each line protects.\n\n```\n{body}\n```",
    "This file changed in a recent deploy. Write the changelog entry and a test "
    "that would fail if the change regressed.\n\n```\n{body}\n```",
]

MAX_LINES = 220


def main() -> int:
    samples = []
    for i, rel in enumerate(TARGETS):
        path = REPO / rel
        if not path.is_file():
            continue
        body = "\n".join(path.read_text(errors="replace").splitlines()[:MAX_LINES])
        samples.append({"prompt": PROMPTS[i % len(PROMPTS)].format(body=body)})
    if not samples:
        print(f"no corpus sources found under {REPO}", file=sys.stderr)
        return 1
    OUT.parent.mkdir(parents=True, exist_ok=True)
    with OUT.open("w") as fh:
        for sample in samples:
            fh.write(json.dumps(sample) + "\n")
    print(f"wrote {len(samples)} prompts to {OUT}")
    return 0


if __name__ == "__main__":
    sys.exit(main())