# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "saturate[hf]",
#     "vllm",
#     "qwen-vl-utils",
# ]
# ///
"""Caption the pre-chunked Prelinger PD mirror with Marlin-2B.

Campaign driver (same validated config as uv-scripts/video's marlin-caption.py,
but consuming the pre-chunked bucket): rows come from /chunks/chunkmeta/*.json —
each chunk carries its ACTUAL stream-copy boundary times, so event timestamps
are offset by recorded starts, not assumed 60s multiples.

Shardable: --shard 0/2 and --shard 1/2 in two parallel jobs write the same
output; saturate's manifest layer keeps them exact and resumable.
"""

import argparse
import json
import re
from pathlib import Path

# All IO via a single rw FUSE mount at /pf (NOT /data — reserved by the platform).
# Local-path output => ParquetSink local writes (rename-free, FUSE-safe) and, crucially,
# existing_ids globs ~5k manifest files at FUSE speed instead of serial HTTP round-trips
# (observed: 60+ min of skip-set loading per job over hf://).
CHUNKS = Path("/pf/chunks")
CMETA = Path("/pf/chunkmeta")
# Fresh output dir: pump's own existing_ids() finds nothing here (instant start).
# Already-done ids are excluded via the consolidated /pf/skip_ids.parquet instead —
# reading 5k manifest sidecars serially costs 60+ min per job on ANY transport.
# Final publish must read BOTH captions/ and captions2/ (or merge them first).
OUTPUT = "/pf/captions2"
SKIP_INDEX = Path("/pf/skip_ids.parquet")
MODEL = "NemoStation/Marlin-2B"

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)
EVENT_LINE = re.compile(r"<(\d+\.?\d*)\s*-\s*(\d+\.?\d*)>\s*(.*)")


def load_rows() -> list[tuple[str, dict]]:
    rows = []
    for mf in sorted(CMETA.glob("*.json")):
        meta = json.loads(mf.read_text())
        ident = meta["identifier"]
        for c in meta["chunks"]:
            path = CHUNKS / ident / c["chunk"]
            rows.append((f"{ident}#{c['start']}", {
                "identifier": ident, "path": str(path),
                "start": c["start"], "end": c["end"],
            }))
    return rows


def parse_events(text: str, offset: float):
    scene, events = "", []
    body = text.split("Events:", 1)
    scene = body[0].replace("Scene:", "", 1).strip()
    for line in (body[1] if len(body) > 1 else "").splitlines():
        m = EVENT_LINE.match(line.strip())
        if m:
            events.append({"start": round(float(m.group(1)) + offset, 2),
                           "end": round(float(m.group(2)) + offset, 2),
                           "text": m.group(3).strip()})
    return scene, events


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shard", default="0/1", help="RANK/WORLD, e.g. 0/2")
    ap.add_argument("--window-max", type=int, default=24)
    args = ap.parse_args()
    rank, world = (int(x) for x in args.shard.split("/"))

    rows = load_rows()
    total = len(rows)
    # Shard CLIENT-SIDE: saturate's pump(shard=) only labels output files, it does
    # not sub-select the input stream (learned at a cost of ~$20 on 2026-07-30).
    rows = rows[rank::world]
    if SKIP_INDEX.exists():
        import pyarrow.parquet as pq
        done = set(pq.read_table(SKIP_INDEX, columns=["id"]).column("id").to_pylist())
        before = len(rows)
        rows = [(rid, r) for rid, r in rows if rid not in done]
        print(f"[plan] skip-index: {len(done)} done ids; shard filtered {before} -> {len(rows)}",
              flush=True)
    print(f"[plan] {total} chunk rows total; shard {rank}/{world} takes {len(rows)}",
          flush=True)

    def to_request(row):
        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, resp):
        text = THINK.sub("", resp["choices"][0]["message"]["content"]).strip()
        usage = resp.get("usage") or {}
        scene, events = parse_events(text, offset=row["start"])
        return {"identifier": row["identifier"], "chunk_start": row["start"],
                "chunk_end": row["end"], "scene": scene,
                "events": json.dumps(events), "caption": text,
                "prompt_tokens": usage.get("prompt_tokens"),
                "completion_tokens": usage.get("completion_tokens")}

    from saturate import Auto, Engine, pump

    with Engine(MODEL, engine="vllm", extra_args=[
            "--hf-overrides", '{"architectures": ["Qwen3_5ForConditionalGeneration"]}',
            "--allowed-local-media-path", str(CHUNKS),
            "--max-model-len", "65536",
            "--mm-processor-cache-gb", "0",
            "--enforce-eager"]) as endpoint:
        stats = pump(rows, to_request=to_request, parse=parse, endpoint=endpoint,
                     output=OUTPUT, window=Auto(initial=8, max_limit=args.window_max),
                     shard=(rank, world))
    print(f"[CAPTION SHARD DONE] {stats}", flush=True)


if __name__ == "__main__":
    main()