# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "saturate[hf]",
#     "vllm",
#     "qwen-vl-utils",
# ]
# ///
"""Saturate + bucket-mount test: Marlin-2B captions Prelinger clips from a FUSE-mounted
bucket, results land as resumable parquet back in the bucket."""

import glob
import os
import re
import time

from saturate import Engine, pump, read_output

CLIPS_DIR = "/clips/clips"
OUTPUT = "hf://buckets/davanstrien/prelinger-sample/captions"

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."
)

THINK = re.compile(r"<think>.*?</think>\s*|^\s*<think>\s*\n*|</think>\s*", re.DOTALL)


def fuse_read_benchmark(paths: list[str]) -> None:
    """Measure sequential read throughput through the FUSE mount."""
    for p in paths:
        t0 = time.perf_counter()
        with open(p, "rb") as f:
            n = len(f.read())
        dt = time.perf_counter() - t0
        print(f"[fuse-bench] {os.path.basename(p)}: {n/1e6:.1f} MB in {dt:.2f}s "
              f"({n/1e6/dt:.1f} MB/s)", flush=True)


def to_request(row: dict) -> dict:
    return {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "video_url", "video_url": {"url": f"file://{row['path']}"}},
                    {"type": "text", "text": CAPTION_PROMPT},
                ],
            }
        ],
        "temperature": 0,
        "max_tokens": 1024,
    }


def parse(row: dict, resp: dict) -> dict:
    raw = resp["choices"][0]["message"]["content"]
    usage = resp.get("usage") or {}
    return {
        "clip": os.path.basename(row["path"]),
        "caption": THINK.sub("", raw).strip(),
        "prompt_tokens": usage.get("prompt_tokens"),
        "completion_tokens": usage.get("completion_tokens"),
    }


def main():
    paths = sorted(glob.glob(f"{CLIPS_DIR}/*.mp4"))
    print(f"found {len(paths)} clips on mount", flush=True)
    assert paths, f"no clips at {CLIPS_DIR} — mount missing?"
    fuse_read_benchmark(paths)

    rows = [(os.path.basename(p), {"path": p}) for p in paths]

    extra_args = [
        "--hf-overrides", '{"architectures": ["Qwen3_5ForConditionalGeneration"]}',
        "--allowed-local-media-path", CLIPS_DIR,
        "--max-model-len", "65536",
        "--enforce-eager",
    ]
    with Engine("NemoStation/Marlin-2B", engine="vllm", extra_args=extra_args) as endpoint:
        stats = pump(
            rows,
            to_request=to_request,
            parse=parse,
            endpoint=endpoint,
            output=OUTPUT,
            flush_every=1,
        )

    print(f"[stats] {stats}", flush=True)
    for rec in read_output(OUTPUT):
        cap = (rec.get("caption") or "")[:200].replace("\n", " ")
        print(f"[row] {rec.get('clip')} err={rec.get('error')} caption={cap}...", flush=True)
    print("SATURATE TEST COMPLETE", flush=True)


if __name__ == "__main__":
    main()