# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "saturate[hf]",
#     "vllm",
#     "qwen-vl-utils",
# ]
# ///
"""Video-delivery benchmark: 4 ways to get an mp4 into a vLLM server, timed.

A: file:// path on a FUSE-mounted bucket (server reads local file)
B: FUSE mount read -> base64 data: URI (bytes travel in request body)
C: fsspec hf://buckets read -> base64 data: URI (HTTP read, no mount)
D: public https URL (server fetches it unauthenticated)

Sequential requests (Fixed(1)) for clean per-request latency.
"""

import base64
import glob
import os
import time

from saturate import Engine, Fixed, pump, bucket_rows

MOUNT_DIR = "/clips/clips"
BUCKET_GLOB = "hf://buckets/davanstrien/prelinger-sample/clips/*.mp4"
PUBLIC_BASE = "https://huggingface.co/datasets/davanstrien/prelinger-clips-sample/resolve/main/clips"
CLIPS = ["AboutBan1935.mp4", "JoanAvoi1947.mp4", "Sleepfor1950.mp4"]

CAPTION_PROMPT = (
    "Provide a spatial description of this clip followed by time-ranged events.\n"
    "For each event, give the time range as <start - end> and a short description."
)


def b64_uri(data: bytes) -> str:
    return "data:video/mp4;base64," + base64.b64encode(data).decode()


def video_request(url: str) -> dict:
    return {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "video_url", "video_url": {"url": url}},
                    {"type": "text", "text": CAPTION_PROMPT},
                ],
            }
        ],
        "temperature": 0,
        "max_tokens": 1024,
    }


def run_mode(name: str, rows, to_request, endpoint) -> dict:
    lat: list[float] = []
    t_req: list[float] = []

    def timed_request(row):
        t_req.append(time.perf_counter())
        return to_request(row)

    def timed_parse(row, resp):
        lat.append(time.perf_counter() - t_req[-1])
        usage = resp.get("usage") or {}
        return {
            "clip": row["clip"],
            "ok": bool(resp["choices"][0]["message"]["content"]),
            "prompt_tokens": usage.get("prompt_tokens"),
        }

    t0 = time.perf_counter()
    stats = pump(
        rows,
        to_request=timed_request,
        parse=timed_parse,
        endpoint=endpoint,
        output=f"/tmp/out/{name}",
        window=Fixed(1),
        flush_every=1,
    )
    wall = time.perf_counter() - t0
    per_req = ", ".join(f"{x:.1f}" for x in lat)
    print(f"[mode {name}] wall={wall:.1f}s rows={stats.rows_processed} "
          f"per-request=[{per_req}]s", flush=True)
    return {"mode": name, "wall": wall, "lat": lat}


def main():
    results = []

    # Source prep, timed separately from inference.
    t0 = time.perf_counter()
    mount_paths = {os.path.basename(p): p for p in sorted(glob.glob(f"{MOUNT_DIR}/*.mp4"))}
    mount_bytes = {n: open(p, "rb").read() for n, p in mount_paths.items()}
    t_mount_read = time.perf_counter() - t0
    print(f"[prep] mount read {len(mount_bytes)} clips in {t_mount_read:.2f}s", flush=True)

    t0 = time.perf_counter()
    fsspec_bytes = {}
    for rid, row in bucket_rows(BUCKET_GLOB, read=True):
        fsspec_bytes[os.path.basename(row["path"])] = row["bytes"]
    t_fsspec_read = time.perf_counter() - t0
    print(f"[prep] fsspec read {len(fsspec_bytes)} clips in {t_fsspec_read:.2f}s", flush=True)

    with Engine(
        "NemoStation/Marlin-2B",
        engine="vllm",
        extra_args=[
            "--hf-overrides", '{"architectures": ["Qwen3_5ForConditionalGeneration"]}',
            "--allowed-local-media-path", MOUNT_DIR,
            "--max-model-len", "65536",
            "--enforce-eager",
        ],
    ) as endpoint:
        # Warm-up request so mode A doesn't pay first-request costs alone.
        warm = [(f"warm-{CLIPS[0]}", {"clip": CLIPS[0]})]
        run_mode("warmup", warm,
                 lambda r: video_request(f"file://{mount_paths[r['clip']]}"), endpoint)

        rows = [(n, {"clip": n}) for n in CLIPS]
        results.append(run_mode(
            "A-file-fuse", rows,
            lambda r: video_request(f"file://{mount_paths[r['clip']]}"), endpoint))
        results.append(run_mode(
            "B-b64-from-mount", rows,
            lambda r: video_request(b64_uri(mount_bytes[r["clip"]])), endpoint))
        results.append(run_mode(
            "C-b64-from-fsspec", rows,
            lambda r: video_request(b64_uri(fsspec_bytes[r["clip"]])), endpoint))
        results.append(run_mode(
            "D-https-public", rows,
            lambda r: video_request(f"{PUBLIC_BASE}/{r['clip']}"), endpoint))

    print("\n===== SUMMARY =====", flush=True)
    print(f"source prep: mount={t_mount_read:.2f}s fsspec={t_fsspec_read:.2f}s", flush=True)
    for r in results:
        mean = sum(r["lat"]) / len(r["lat"]) if r["lat"] else float("nan")
        print(f"{r['mode']:20} wall={r['wall']:6.1f}s mean-request={mean:5.1f}s", flush=True)
    print("DELIVERY BENCH COMPLETE", flush=True)


if __name__ == "__main__":
    main()