# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "saturate[hf]",
#     "vllm",
#     "qwen-vl-utils",
#     "opencv-python-headless",
# ]
# ///
"""Extended video bench: IA-direct fetch + long-video failure mode, fps sweep,
frames-as-images vs true video. Talks to the vLLM server via saturate's transport
only where pump adds value; single probes use plain httpx for clarity."""

import base64
import glob
import re
import time

import httpx

from saturate import Engine

MOUNT_DIR = "/clips/clips"
IA_FULL_FILM = "https://archive.org/download/AboutBan1935/AboutBan1935_512kb.mp4"
PROBE_CLIP = "AboutBan1935.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 chat(endpoint: str, content: list, timeout: float = 900, **kwargs) -> dict:
    body = {"model": "NemoStation/Marlin-2B",
            "messages": [{"role": "user", "content": content}],
            "temperature": 0, "max_tokens": 1024, **kwargs}
    r = httpx.post(f"{endpoint}/chat/completions", json=body, timeout=timeout)
    return {"status": r.status_code, "json": r.json() if r.status_code == 200 else r.text[:500]}


def video_part(url: str) -> dict:
    return {"type": "video_url", "video_url": {"url": url}}


def report(tag: str, t0: float, out: dict) -> None:
    dt = time.perf_counter() - t0
    if out["status"] == 200:
        j = out["json"]
        usage = j.get("usage") or {}
        text = THINK.sub("", j["choices"][0]["message"]["content"]).strip()
        events = len(re.findall(r"<\d+\.?\d*\s*-\s*\d+\.?\d*>", text))
        print(f"[{tag}] ok {dt:.1f}s prompt_tokens={usage.get('prompt_tokens')} "
              f"events={events}", flush=True)
        print(f"[{tag}] first-300: {text[:300].replace(chr(10), ' | ')}", flush=True)
    else:
        print(f"[{tag}] HTTP {out['status']} in {dt:.1f}s: {out['json'][:300]}", flush=True)


def sample_frames_b64(path: str, n: int = 16) -> list:
    import cv2
    cap = cv2.VideoCapture(path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    parts = []
    for i in range(n):
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(i * total / n))
        ok, frame = cap.read()
        if not ok:
            continue
        ok, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85])
        parts.append({"type": "image_url", "image_url": {
            "url": "data:image/jpeg;base64," + base64.b64encode(buf).decode()}})
    cap.release()
    print(f"[frames] sampled {len(parts)} jpegs, "
          f"{sum(len(p['image_url']['url']) for p in parts)/1e6:.1f} MB total", flush=True)
    return parts


def main():
    clip = f"{MOUNT_DIR}/{PROBE_CLIP}"
    assert glob.glob(clip), f"missing {clip}"

    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
        t0 = time.perf_counter()
        report("warmup", t0, chat(endpoint, [video_part(f"file://{clip}"),
                                             {"type": "text", "text": CAPTION_PROMPT}]))

        # 1) fps sweep on the 60s clip (token cost vs timestamp behavior)
        for fps in (0.5, 1.0, 2.0):
            t0 = time.perf_counter()
            out = chat(endpoint, [video_part(f"file://{clip}"),
                                  {"type": "text", "text": CAPTION_PROMPT}],
                       mm_processor_kwargs={"fps": fps})
            report(f"fps={fps}", t0, out)

        # 2) frames-as-images vs true video (16 uniform jpegs, no temporal encoding)
        frames = sample_frames_b64(clip, n=16)
        t0 = time.perf_counter()
        out = chat(endpoint, frames + [{"type": "text", "text": CAPTION_PROMPT}])
        report("frames-as-images", t0, out)

        # 3) IA-direct: server fetches the FULL film straight from archive.org.
        #    Expect either success (if frame cap kicks in) or a clean context-length 400.
        t0 = time.perf_counter()
        out = chat(endpoint, [video_part(IA_FULL_FILM),
                              {"type": "text", "text": CAPTION_PROMPT}], timeout=1200)
        report("ia-direct-full-film", t0, out)

        # 3b) IA-direct with aggressive sampling caps, if raw failed or succeeded oddly
        t0 = time.perf_counter()
        out = chat(endpoint, [video_part(IA_FULL_FILM),
                              {"type": "text", "text": CAPTION_PROMPT}], timeout=1200,
                   mm_processor_kwargs={"fps": 0.2})
        report("ia-direct-fps0.2", t0, out)

    print("EXTENDED BENCH COMPLETE", flush=True)


if __name__ == "__main__":
    main()