#!/usr/bin/env python3
"""Issue one non-streaming completion request and record wall/usage."""

from __future__ import annotations

import argparse
import json
import time
import urllib.request
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, required=True)
    parser.add_argument("--request", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--timeout-seconds", type=float, default=1800)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    body = json.loads(args.request.read_text())
    expected_input = len(body["prompt"])
    expected_output = int(body["max_tokens"])
    request = urllib.request.Request(
        f"http://127.0.0.1:{args.port}/v1/completions",
        data=json.dumps(body, separators=(",", ":")).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.monotonic()
    with urllib.request.urlopen(
        request, timeout=args.timeout_seconds
    ) as response:
        payload = json.load(response)
    wall_ms = (time.monotonic() - started) * 1000
    usage = payload.get("usage") or {}
    observed = (
        int(usage.get("prompt_tokens", -1)),
        int(usage.get("completion_tokens", -1)),
    )
    if observed != (expected_input, expected_output):
        raise ValueError(
            f"usage mismatch: expected {expected_input}+{expected_output}, "
            f"observed {observed[0]}+{observed[1]}"
        )
    result = {
        "wall_ms": wall_ms,
        "input_tokens": observed[0],
        "output_tokens": observed[1],
        "usage": usage,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(json.dumps(result, sort_keys=True), flush=True)


if __name__ == "__main__":
    main()