# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "saturate[hf]",
#     "vllm",
#     "qwen-vl-utils",
#     "psutil",
# ]
# ///
"""Stress test: 33 distinct 60s clips through Marlin-2B on vLLM with saturate's
DEFAULT Auto window (nothing set). Question: does the controller overshoot on a
video endpoint whose per-request cost is frontend CPU decode, and does the server
survive? RSS sampler watches total memory the whole time.

Ends with one IA-direct full-film request, printing the complete event list
(timestamp-span check the extended bench truncated away)."""

import glob
import os
import re
import threading
import time

import httpx
import psutil

from saturate import Auto, Engine, pump

CLIPS_DIR = "/clips/stress_clips"
OUTPUT = "hf://buckets/davanstrien/prelinger-sample/captions-stress"
IA_FULL_FILM = "https://archive.org/download/AboutBan1935/AboutBan1935_512kb.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."
)

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


def rss_sampler(stop: threading.Event) -> None:
    """Every 5s: total used RAM + biggest process, so OOM pressure is visible live."""
    peak = 0.0
    while not stop.is_set():
        vm = psutil.virtual_memory()
        used_gb = (vm.total - vm.available) / 1e9
        peak = max(peak, used_gb)
        procs = sorted(psutil.process_iter(["name", "memory_info"]),
                       key=lambda p: p.info["memory_info"].rss if p.info["memory_info"] else 0,
                       reverse=True)
        top = procs[0]
        print(f"[rss] used={used_gb:.1f}G/{vm.total/1e9:.0f}G peak={peak:.1f}G "
              f"top={top.info['name']}:{top.info['memory_info'].rss/1e9:.1f}G", flush=True)
        stop.wait(5)


def main():
    paths = sorted(glob.glob(f"{CLIPS_DIR}/*.mp4"))
    print(f"found {len(paths)} distinct clips", flush=True)
    assert len(paths) >= 30, f"expected >=30 clips, got {len(paths)}"

    stop = threading.Event()
    threading.Thread(target=rss_sampler, args=(stop,), daemon=True).start()

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

    with Engine(
        "NemoStation/Marlin-2B",
        engine="vllm",
        extra_args=[
            "--hf-overrides", '{"architectures": ["Qwen3_5ForConditionalGeneration"]}',
            "--allowed-local-media-path", CLIPS_DIR,
            "--max-model-len", "65536",
            "--enforce-eager",
            "--mm-processor-cache-gb", "0",
        ],
    ) as endpoint:
        t0 = time.perf_counter()
        stats = pump(
            rows,
            to_request=lambda r: {
                "messages": [{"role": "user", "content": [
                    {"type": "video_url", "video_url": {"url": f"file://{r['path']}"}},
                    {"type": "text", "text": CAPTION_PROMPT},
                ]}],
                "temperature": 0, "max_tokens": 1024,
            },
            parse=lambda r, resp: {
                "clip": os.path.basename(r["path"]),
                "caption": THINK.sub("", resp["choices"][0]["message"]["content"]).strip(),
                "prompt_tokens": (resp.get("usage") or {}).get("prompt_tokens"),
            },
            endpoint=endpoint,
            output=OUTPUT,
            # positive control: the capped window we recommend for video recipes
            window=Auto(initial=8, max_limit=24),
        )
        wall = time.perf_counter() - t0
        print(f"[stress-stats] {stats}", flush=True)
        print(f"[stress-summary] {stats.rows_processed} ok / {stats.rows_failed} failed "
              f"in {wall:.0f}s = {wall/max(1, stats.rows_processed):.1f}s/clip effective; "
              f"final_limit={stats.final_limit} breaker_opens={stats.breaker_opens}", flush=True)

    stop.set()
    print("STRESS TEST COMPLETE", flush=True)


if __name__ == "__main__":
    main()