"""Observe the original upstream benchmark without changing its sampling/timing.

Run on the approved single-GPU Pod, after downloading the pinned model locally.
The wrapper records actual requests, completion and sampled device memory.
"""
import argparse
import hashlib
import importlib.metadata
import json
import math
import subprocess
import sys
import threading
import time
from pathlib import Path

import torch
from vllm.benchmarks import throughput
from vllm.utils.argparse_utils import FlexibleArgumentParser


def digest(value):
    return hashlib.sha256(json.dumps(value, ensure_ascii=False, sort_keys=True,
                                     separators=(",", ":")).encode()).hexdigest()


def gpu_snapshot():
    output = subprocess.check_output([
        "nvidia-smi", "--query-gpu=uuid,name,driver_version,memory.total,memory.used",
        "--format=csv,noheader,nounits"], text=True, timeout=3)
    rows = output.strip().splitlines()
    if len(rows) != 1:
        raise RuntimeError("exactly one physical GPU required")
    uuid, name, driver, total, used = [x.strip() for x in rows[0].split(",")]
    return {"uuid": uuid, "name": name, "driver": driver,
            "memory_total_mib": int(total), "memory_used_mib": int(used)}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("observation", type=Path)
    parser.add_argument("--expected-version", required=True)
    wrapper, arguments = parser.parse_known_args()
    record = {"expected_version": wrapper.expected_version,
              "vllm_version": importlib.metadata.version("vllm"),
              "cuda_runtime": torch.version.cuda, "torch": torch.__version__,
              "gpu_before": gpu_snapshot(), "generated_fixtures": [],
              "measured_runs": 0,
              "execution": "NOT_RUN"}
    actual_version = record["vllm_version"].split("+")[0]
    if actual_version != wrapper.expected_version:
        raise RuntimeError("installed vLLM does not match selected release")
    if not torch.cuda.is_available() or torch.cuda.device_count() != 1:
        raise RuntimeError("one CUDA GPU required")
    original_requests = throughput.get_requests
    original_run = throughput.run_vllm

    def observe_requests(args, tokenizer):
        requests = original_requests(args, tokenizer)
        data = [{"prompt": r.prompt, "prompt_len": r.prompt_len,
                 "expected_output_len": r.expected_output_len} for r in requests]
        record["generated_fixtures"].append({"seed": args.seed,
            "requests": len(data), "sha256": digest(data), "data": data})
        return requests

    def observe_run(*args, **kwargs):
        elapsed, outputs = original_run(*args, **kwargs)
        if outputs is None or len(outputs) != 32:
            raise RuntimeError("expected exactly 32 actual RequestOutput objects")
        if any(not output.finished or len(output.outputs) != 1
               or len(output.outputs[0].token_ids) != 32 for output in outputs):
            raise RuntimeError("all 32 requests must finish with 32 output tokens")
        input_tokens = [output.prompt_token_ids for output in outputs]
        if any(tokens is None for tokens in input_tokens):
            raise RuntimeError("actual input tokens are unavailable")
        if not math.isfinite(elapsed) or elapsed <= 0:
            raise RuntimeError("measured elapsed time must be finite and positive")
        record.update(completed_requests=len(outputs),
                      actual_input_tokens_sha256=digest(input_tokens),
                      actual_input_tokens=input_tokens,
                      upstream_elapsed_seconds=elapsed)
        record["measured_runs"] += 1
        return elapsed, outputs

    throughput.get_requests = observe_requests
    throughput.run_vllm = observe_run
    samples, sampling_errors = [], []
    stop = threading.Event()

    def sample_memory():
        while not stop.is_set():
            try:
                samples.append({"at_unix": time.time(), **gpu_snapshot()})
            except Exception as exc:
                sampling_errors.append(str(exc))
            stop.wait(1)

    monitor = threading.Thread(target=sample_memory, daemon=True)
    monitor.start()
    try:
        benchmark_parser = FlexibleArgumentParser()
        throughput.add_cli_args(benchmark_parser)
        throughput.main(benchmark_parser.parse_args(arguments))
        fixtures = record["generated_fixtures"]
        if [(fixture["seed"], fixture["requests"]) for fixture in fixtures] != [(1, 8), (0, 32)]:
            raise RuntimeError("expected complete warmup8/seed1 and measured32/seed0 fixtures")
        if any(request["expected_output_len"] != 32
               for fixture in fixtures for request in fixture["data"]):
            raise RuntimeError("all fixture requests must specify 32 output tokens")
        if record["measured_runs"] != 1 or record.get("completed_requests") != 32:
            raise RuntimeError("exactly one complete measured run required")
        record["execution"] = "PASS"
    except Exception as exc:
        record["execution"] = "INFRA_OR_WORKLOAD_ERROR"
        record["error"] = str(exc)
        raise
    finally:
        stop.set()
        monitor.join(timeout=4)
        record.update(device_memory_samples=samples, memory_sampling_errors=sampling_errors,
                      sampled_device_memory_peak_mib=max(
                          (sample["memory_used_mib"] for sample in samples), default=None),
                      memory_measurement_scope="1-second device-wide samples; not allocator peak",
                      finished_at_unix=time.time())
        wrapper.observation.parent.mkdir(parents=True, exist_ok=True)
        wrapper.observation.write_text(json.dumps(record, indent=2) + "\n")


if __name__ == "__main__":
    main()