scripts/prepare_streamingbench.py

#!/usr/bin/env python3
"""Prepare StreamingBench for the StreamDelta streaming inference runner."""

from __future__ import annotations

import argparse
import ast
import collections
import concurrent.futures
import csv
import hashlib
import json
import math
import os
import re
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple


EXPECTED_CSV_ROWS = {
    "Real_Time_Visual_Understanding.csv": 2500,
    "Sequential_Question_Answering.csv": 250,
    "Contextual_Understanding.csv": 500,
    "Omni_Source_Understanding.csv": 1000,
    "Proactive_Output.csv": 250,
}
EXPECTED_SUITE_COUNTS = {
    "Real-Time Visual Understanding": 2500,
    "Sequential Question Answering": 250,
    "Contextual Understanding": 500,
    "Omni-Source Understanding": 1000,
    "Proactive Output": 250,
}
EXPECTED_VIDEO_COUNTS = {
    "Real-Time Visual Understanding": 500,
    "Sequential Question Answering": 50,
    "Anomaly Context Understanding": 50,
    "Misleading Context Understanding": 50,
    "Emotion Recognition": 50,
    "Multimodal Alignment": 50,
    "Scene Understanding": 50,
    "Source Discrimination": 50,
    "Proactive Output": 50,
}
CSV_SPECS = (
    (
        "Real_Time_Visual_Understanding.csv",
        "Real-Time Visual Understanding",
    ),
    (
        "Sequential_Question_Answering.csv",
        "Sequential Question Answering",
    ),
    ("Contextual_Understanding.csv", "Contextual Understanding"),
    ("Omni_Source_Understanding.csv", "Omni-Source Understanding"),
    ("Proactive_Output.csv", "Proactive Output"),
)
TASK_TO_MEDIA_DIRECTORY = {
    "Misleading Context Recognition": "Misleading Context Understanding",
    "Anomaly Context Understanding": "Anomaly Context Understanding",
    "Emotion Recognition": "Emotion Recognition",
    "Multimodal Alignment": "Multimodal Alignment",
    "Scene Understanding": "Scene Understanding",
    "Source Discrimination": "Source Discrimination",
}
SUITE_TO_MEDIA_DIRECTORY = {
    "Real-Time Visual Understanding": "Real-Time Visual Understanding",
    "Sequential Question Answering": "Sequential Question Answering",
    "Proactive Output": "Proactive Output",
}
MEDIA_DIRECTORY_SLUGS = {
    name: re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_")
    for name in EXPECTED_VIDEO_COUNTS
}
QUESTION_ID_RE = re.compile(r"_sample_(\d+)_([1-5])$")


def write_json_atomic(path: Path, value: Any, *, pretty: bool = False) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        "w",
        encoding="utf-8",
        dir=path.parent,
        prefix=f".{path.name}.",
        delete=False,
    ) as handle:
        json.dump(
            value,
            handle,
            ensure_ascii=False,
            indent=2 if pretty else None,
        )
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def read_json(path: Path) -> Any:
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def parse_timestamp(value: str) -> int:
    parts = str(value).strip().split(":")
    if len(parts) == 2:
        hours = 0
        minutes = int(parts[0])
        seconds = float(parts[1])
    elif len(parts) == 3:
        hours, minutes = map(int, parts[:2])
        seconds = float(parts[2])
    else:
        raise ValueError(f"invalid StreamingBench timestamp: {value!r}")
    if hours < 0 or not 0 <= minutes < 60 or not 0 <= seconds < 60:
        raise ValueError(f"invalid StreamingBench timestamp: {value!r}")
    return int(math.floor(hours * 3600 + minutes * 60 + seconds))


def normalize_options(raw: str) -> List[str]:
    options = ast.literal_eval(raw)
    if not isinstance(options, list) or not options:
        raise ValueError(f"invalid options list: {raw!r}")
    normalized: List[str] = []
    for index, value in enumerate(options):
        text = str(value).strip()
        expected = chr(ord("A") + index)
        text = re.sub(
            rf"^\s*{re.escape(expected)}\s*[.)]\s*",
            "",
            text,
            count=1,
            flags=re.IGNORECASE,
        )
        normalized.append(text)
    return normalized


def parse_source_number(question_id: str) -> int:
    match = QUESTION_ID_RE.search(question_id)
    if not match:
        raise ValueError(f"cannot map question_id to source video: {question_id}")
    return int(match.group(1))


def media_directory_for(suite: str, task_type: str) -> str:
    if suite in SUITE_TO_MEDIA_DIRECTORY:
        return SUITE_TO_MEDIA_DIRECTORY[suite]
    try:
        return TASK_TO_MEDIA_DIRECTORY[task_type]
    except KeyError as exc:
        raise ValueError(
            f"no media mapping for suite={suite!r}, task_type={task_type!r}"
        ) from exc


def find_source_video(source_root: Path, directory: str, number: int) -> Path:
    sample_dir = source_root / directory / f"sample_{number}"
    canonical = sample_dir / "video.mp4"
    if canonical.is_file() and canonical.stat().st_size:
        return canonical
    candidates = sorted(
        path
        for path in sample_dir.glob("*.mp4")
        if path.is_file() and path.stat().st_size
    )
    if len(candidates) == 1:
        return candidates[0]
    raise FileNotFoundError(
        f"expected one source video under {sample_dir}, found {candidates}"
    )


def source_key(directory: str, number: int) -> str:
    return f"{MEDIA_DIRECTORY_SLUGS[directory]}/sample_{number}"


def frame_marker(frame_dir: Path) -> Path:
    return (
        frame_dir.parents[1]
        / ".streamdelta_complete"
        / frame_dir.parent.name
        / f"{frame_dir.name}.json"
    )


def completed_frame_count(
    frame_dir: Path,
    source: Path,
    duration: int,
) -> int:
    marker = frame_marker(frame_dir)
    first = frame_dir / "000000.jpg"
    if not marker.is_file() or not first.is_file():
        return 0
    try:
        payload = read_json(marker)
        count = int(payload["frames"])
    except (KeyError, TypeError, ValueError, json.JSONDecodeError):
        return 0
    stat = source.stat()
    if (
        count < 1
        or not (frame_dir / f"{count - 1:06d}.jpg").is_file()
        or int(payload.get("duration", -1)) != duration
        or int(payload.get("source_size", -1)) != stat.st_size
        or int(payload.get("source_mtime_ns", -1)) != stat.st_mtime_ns
    ):
        return 0
    return count


def resolve_ffmpeg(explicit: str | None) -> str:
    if explicit:
        resolved = shutil.which(explicit) or explicit
        if Path(resolved).is_file():
            return resolved
        raise FileNotFoundError(f"ffmpeg not found: {explicit}")
    system = shutil.which("ffmpeg")
    if system:
        return system
    try:
        import imageio_ffmpeg

        return imageio_ffmpeg.get_ffmpeg_exe()
    except (ImportError, RuntimeError) as exc:
        raise RuntimeError(
            "ffmpeg is required; install ffmpeg or imageio-ffmpeg"
        ) from exc


def resolve_ffprobe(ffmpeg: str) -> str:
    system = shutil.which("ffprobe")
    if system:
        return system
    sibling = Path(ffmpeg).with_name("ffprobe")
    if sibling.is_file():
        return str(sibling)
    raise FileNotFoundError("ffprobe is required to validate video durations")


def probe_duration(ffprobe: str, source: Path) -> float:
    result = subprocess.run(
        [
            ffprobe,
            "-v",
            "error",
            "-show_entries",
            "format=duration",
            "-of",
            "default=noprint_wrappers=1:nokey=1",
            str(source),
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
    )
    if result.returncode:
        raise RuntimeError(
            f"ffprobe failed for {source}: {result.stderr[-2000:].strip()}"
        )
    duration = float(result.stdout.strip())
    if not math.isfinite(duration) or duration <= 0:
        raise ValueError(f"invalid video duration for {source}: {duration}")
    return duration


def extract_frames(
    *,
    ffmpeg: str,
    source: Path,
    frame_dir: Path,
    duration: int,
    force: bool,
) -> int:
    existing = completed_frame_count(frame_dir, source, duration)
    if existing and not force:
        return existing
    frame_dir.parent.mkdir(parents=True, exist_ok=True)
    temporary = frame_dir.with_name(f".{frame_dir.name}.{os.getpid()}.tmp")
    if temporary.exists():
        shutil.rmtree(temporary)
    temporary.mkdir(parents=True)
    command = [
        ffmpeg,
        "-hide_banner",
        "-loglevel",
        "error",
        "-y",
        "-i",
        str(source),
        "-t",
        str(duration),
        "-vf",
        "fps=1",
        "-start_number",
        "0",
        "-q:v",
        "2",
        str(temporary / "%06d.jpg"),
    ]
    try:
        result = subprocess.run(
            command,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            text=True,
        )
        if result.returncode:
            raise RuntimeError(
                f"frame extraction failed for {source}: "
                f"{result.stderr[-2000:].strip()}"
            )
        frames = sorted(temporary.glob("*.jpg"))
        if not frames or frames[0].name != "000000.jpg":
            raise RuntimeError(f"no 1 fps frames generated for {source}")
        if frame_dir.exists():
            shutil.rmtree(frame_dir)
        os.replace(temporary, frame_dir)
        stat = source.stat()
        write_json_atomic(
            frame_marker(frame_dir),
            {
                "frames": len(frames),
                "duration": duration,
                "source": str(source.resolve()),
                "source_size": stat.st_size,
                "source_mtime_ns": stat.st_mtime_ns,
            },
            pretty=True,
        )
        return len(frames)
    finally:
        if temporary.exists():
            shutil.rmtree(temporary)


def replace_symlink(link: Path, target: Path) -> None:
    link.parent.mkdir(parents=True, exist_ok=True)
    relative = os.path.relpath(target.resolve(), link.parent.resolve())
    if link.is_symlink() and os.readlink(link) == relative and link.exists():
        return
    if link.exists() or link.is_symlink():
        if link.is_dir() and not link.is_symlink():
            raise IsADirectoryError(f"refusing to replace directory: {link}")
        link.unlink()
    temporary = link.with_name(f".{link.name}.{os.getpid()}.tmp")
    temporary.unlink(missing_ok=True)
    os.symlink(relative, temporary)
    os.replace(temporary, link)


def load_csv_rows(annotation_root: Path) -> List[Tuple[str, Dict[str, str]]]:
    output: List[Tuple[str, Dict[str, str]]] = []
    for filename, suite in CSV_SPECS:
        path = annotation_root / filename
        if not path.is_file():
            raise FileNotFoundError(f"missing annotation CSV: {path}")
        with path.open(encoding="utf-8-sig", newline="") as handle:
            rows = list(csv.DictReader(handle))
        expected = EXPECTED_CSV_ROWS[filename]
        if len(rows) != expected:
            raise ValueError(
                f"{filename}: expected {expected} rows, got {len(rows)}"
            )
        output.extend((suite, row) for row in rows)
    return output


def build_sequential_context(
    grouped_rows: Dict[Tuple[str, int], List[Dict[str, str]]],
) -> Dict[str, List[Dict[str, str]]]:
    context: Dict[str, List[Dict[str, str]]] = {}
    for rows in grouped_rows.values():
        ordered = sorted(rows, key=lambda row: parse_timestamp(row["time_stamp"]))
        history: List[Dict[str, str]] = []
        for row in ordered:
            context[row["question_id"]] = [dict(item) for item in history]
            history.append({
                "question": row["question"].strip(),
                "time_stamp": row["time_stamp"].strip(),
            })
    return context


def prepare_samples(
    source_root: Path,
    output_root: Path,
    proactive_tail_seconds: int,
    *,
    ffprobe: str,
    probe_workers: int,
) -> Tuple[
    List[Dict[str, Any]],
    Dict[str, Dict[str, Any]],
    Dict[str, int],
]:
    annotation_root = source_root / "StreamingBench"
    raw_rows = load_csv_rows(annotation_root)
    sequential_groups: Dict[Tuple[str, int], List[Dict[str, str]]] = (
        collections.defaultdict(list)
    )
    for suite, row in raw_rows:
        if suite == "Sequential Question Answering":
            sequential_groups[(suite, parse_source_number(row["question_id"]))].append(
                row
            )
    sequential_context = build_sequential_context(sequential_groups)

    pending: List[Tuple[str, Dict[str, str], str, int, Path, str, Path]] = []
    source_paths: Dict[str, Path] = {}
    for suite, raw in raw_rows:
        question_id = raw["question_id"].strip()
        task_type = raw["task_type"].strip()
        number = parse_source_number(question_id)
        media_directory = media_directory_for(suite, task_type)
        source = find_source_video(source_root, media_directory, number)
        key = source_key(media_directory, number)
        frame_dir = output_root / "frames" / key
        pending.append((
            suite,
            raw,
            task_type,
            number,
            source,
            key,
            frame_dir,
        ))
        source_paths[key] = source

    source_durations: Dict[str, float] = {}
    with concurrent.futures.ThreadPoolExecutor(
        max_workers=max(1, probe_workers)
    ) as pool:
        futures = {
            pool.submit(probe_duration, ffprobe, path): key
            for key, path in source_paths.items()
        }
        for future in concurrent.futures.as_completed(futures):
            key = futures[future]
            source_durations[key] = float(future.result())

    samples: List[Dict[str, Any]] = []
    sources: Dict[str, Dict[str, Any]] = {}
    required_seconds: Dict[str, int] = collections.defaultdict(int)

    for suite, raw, task_type, number, source, key, frame_dir in pending:
        question_id = raw["question_id"].strip()
        media_directory = media_directory_for(suite, task_type)
        source_duration = source_durations[key]
        available_last_second = max(0, int(math.ceil(source_duration)) - 1)
        annotated_question_time = parse_timestamp(raw["time_stamp"])
        question_time = min(annotated_question_time, available_last_second)
        sources[key] = {
            "source_id": key,
            "media_directory": media_directory,
            "sample_number": number,
            "video": str(source.resolve()),
            "frames": str(frame_dir.resolve()),
            "duration_seconds": source_duration,
            "available_last_second": available_last_second,
        }

        base: Dict[str, Any] = {
            "id": question_id,
            "benchmark": "streamingbench",
            "suite": suite,
            "task": task_type,
            "task_type": task_type,
            "question": raw["question"].strip(),
            "realtime": question_time,
            "annotated_realtime": annotated_question_time,
            "question_time_clamped": question_time != annotated_question_time,
            "time_stamp": raw["time_stamp"].strip(),
            "temporal_clue_type": raw["temporal_clue_type"].strip(),
            "frames_required": raw["frames_required"].strip(),
            "source_id": key,
            "source_video": str(source.resolve()),
            "video": str(frame_dir.resolve()),
        }
        if suite == "Proactive Output":
            annotated_target_time = parse_timestamp(
                raw["ground_truth_time_stamp"]
            )
            target_time = min(annotated_target_time, available_last_second)
            base.update({
                "proactive_output": True,
                "target_time": target_time,
                "annotated_target_time": annotated_target_time,
                "target_time_clamped": target_time != annotated_target_time,
                "ground_truth_time_stamp": raw[
                    "ground_truth_time_stamp"
                ].strip(),
                "expected_output": raw["ground_truth_output"].strip(),
                "gt": raw["ground_truth_output"].strip(),
            })
            base["stream_end_time"] = min(
                available_last_second + 1,
                target_time + proactive_tail_seconds + 1,
            )
            required_seconds[key] = max(
                required_seconds[key],
                int(base["stream_end_time"]),
            )
        else:
            options = normalize_options(raw["options"])
            answer = raw["answer"].strip().upper()
            if not re.fullmatch(r"[A-Z]", answer):
                raise ValueError(
                    f"{question_id}: invalid answer letter {answer!r}"
                )
            base.update({
                "options": options,
                "gt": answer,
                "stream_end_time": question_time + 1,
            })
            required_seconds[key] = max(
                required_seconds[key],
                question_time + 1,
            )
            if suite == "Sequential Question Answering":
                history = sequential_context[question_id]
                base["native_question"] = base["question"]
                base["sequential_context"] = history
                if history:
                    rendered = "\n".join(
                        (
                            f"- At {item['time_stamp']}, Q: "
                            f"{item['question']}"
                        )
                        for item in history
                    )
                    base["question"] = (
                        "Earlier questions from this same stream are listed "
                        "below only to resolve references such as 'the person "
                        "mentioned in the previous question'. Their answers "
                        "are not provided; infer the referents from the video "
                        "and MEMORY.\n"
                        f"{rendered}\n\n"
                        f"Current question: {base['native_question']}"
                    )
        samples.append(base)

    validate_samples(samples, sources)
    return samples, sources, dict(required_seconds)


def validate_samples(
    samples: Sequence[Dict[str, Any]],
    sources: Dict[str, Dict[str, Any]],
) -> None:
    ids = [str(sample["id"]) for sample in samples]
    if len(samples) != 4500 or len(set(ids)) != 4500:
        raise ValueError(
            f"expected 4500 unique QA rows, got rows={len(samples)}, "
            f"unique={len(set(ids))}"
        )
    suite_counts = collections.Counter(sample["suite"] for sample in samples)
    if dict(suite_counts) != EXPECTED_SUITE_COUNTS:
        raise ValueError(
            f"unexpected suite counts: {dict(suite_counts)}; "
            f"expected {EXPECTED_SUITE_COUNTS}"
        )
    source_counts = collections.Counter(
        source["media_directory"] for source in sources.values()
    )
    if dict(source_counts) != EXPECTED_VIDEO_COUNTS:
        raise ValueError(
            f"unexpected video counts: {dict(source_counts)}; "
            f"expected {EXPECTED_VIDEO_COUNTS}"
        )


def run_frame_jobs(
    jobs: Iterable[Tuple[str, Any]],
    workers: int,
) -> Dict[str, int]:
    job_list = list(jobs)
    results: Dict[str, int] = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {
            pool.submit(function): name for name, function in job_list
        }
        completed = 0
        for future in concurrent.futures.as_completed(futures):
            name = futures[future]
            results[name] = int(future.result())
            completed += 1
            if completed == 1 or completed % 25 == 0 or completed == len(job_list):
                print(
                    f"[prepare] frames completed {completed}/{len(job_list)}: "
                    f"{name}"
                )
    return results


def reconcile_sample_times(
    samples: Sequence[Dict[str, Any]],
    frame_counts: Dict[str, int],
    proactive_tail_seconds: int,
) -> None:
    for sample in samples:
        sample_id = str(sample["id"])
        source_id = str(sample["source_id"])
        frame_count = int(frame_counts[source_id])
        if frame_count < 1:
            raise ValueError(
                f"{sample_id}: source {source_id} produced no frames"
            )
        last_frame = frame_count - 1
        annotated_question_time = int(sample["annotated_realtime"])
        sample["realtime"] = min(annotated_question_time, last_frame)
        sample["question_time_clamped"] = (
            sample["realtime"] != annotated_question_time
        )
        if sample.get("proactive_output"):
            annotated_target_time = int(sample["annotated_target_time"])
            sample["target_time"] = min(annotated_target_time, last_frame)
            sample["target_time_clamped"] = (
                sample["target_time"] != annotated_target_time
            )
            sample["stream_end_time"] = min(
                frame_count,
                int(sample["target_time"]) + proactive_tail_seconds + 1,
            )
        else:
            sample["stream_end_time"] = int(sample["realtime"]) + 1


def audit_prepared(output_root: Path, dataset_path: Path) -> Dict[str, Any]:
    samples = read_json(dataset_path)
    sources_path = output_root / "sources.json"
    sources = read_json(sources_path)
    validate_samples(samples, sources)
    missing_frames: List[str] = []
    missing_audio: List[str] = []
    invalid_times: List[str] = []
    source_ids_with_frames: set[str] = set()
    frame_counts: Dict[str, int] = {}
    for sample in samples:
        frame_dir = Path(sample["video"])
        frame_key = str(frame_dir)
        if frame_key not in frame_counts:
            marker = frame_marker(frame_dir)
            try:
                frame_counts[frame_key] = int(read_json(marker)["frames"])
            except (
                FileNotFoundError,
                KeyError,
                TypeError,
                ValueError,
                json.JSONDecodeError,
            ):
                frame_counts[frame_key] = (
                    len(list(frame_dir.glob("*.jpg")))
                    if frame_dir.is_dir()
                    else 0
                )
        frame_count = frame_counts[frame_key]
        if not frame_count:
            missing_frames.append(str(sample["id"]))
        else:
            source_ids_with_frames.add(str(sample["source_id"]))
            stream_end_time = int(sample.get("stream_end_time") or 0)
            if stream_end_time < 1 or frame_count < stream_end_time:
                invalid_times.append(str(sample["id"]))
        audio = output_root / "audio_videos" / f"{sample['id']}.mp4"
        if not audio.is_file():
            missing_audio.append(str(sample["id"]))
    report = {
        "dataset": str(dataset_path.resolve()),
        "rows": len(samples),
        "unique_ids": len({str(sample["id"]) for sample in samples}),
        "source_videos": len(sources),
        "source_videos_with_frames": len(source_ids_with_frames),
        "suite_counts": dict(collections.Counter(
            sample["suite"] for sample in samples
        )),
        "task_counts": dict(collections.Counter(
            sample["task"] for sample in samples
        )),
        "clamped_question_times": [
            {
                "id": sample["id"],
                "annotated": sample["annotated_realtime"],
                "effective": sample["realtime"],
            }
            for sample in samples
            if sample.get("question_time_clamped")
        ],
        "clamped_proactive_target_times": [
            {
                "id": sample["id"],
                "annotated": sample["annotated_target_time"],
                "effective": sample["target_time"],
            }
            for sample in samples
            if sample.get("target_time_clamped")
        ],
        "missing_frame_samples": missing_frames,
        "missing_audio_videos": missing_audio,
        "insufficient_frame_samples": invalid_times,
        "passed": not missing_frames and not missing_audio and not invalid_times,
    }
    if not report["passed"]:
        raise RuntimeError(
            "StreamingBench media audit failed: "
            f"missing_frames={len(missing_frames)}, "
            f"missing_audio={len(missing_audio)}, "
            f"insufficient_frames={len(invalid_times)}"
        )
    return report


def main() -> None:
    repository = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--source-root",
        type=Path,
        default=Path("data/StreamingBench-source"),
    )
    parser.add_argument(
        "--output-root",
        type=Path,
        default=repository / "data" / "StreamingBench",
    )
    parser.add_argument("--jobs", type=int, default=8)
    parser.add_argument("--ffmpeg")
    parser.add_argument("--proactive-tail-seconds", type=int, default=10)
    parser.add_argument("--skip-frames", action="store_true")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--check-only", action="store_true")
    args = parser.parse_args()

    source_root = args.source_root.resolve()
    output_root = args.output_root.resolve()
    dataset_path = output_root / "streamingbench_streamdelta_eval.json"
    audit_path = output_root / "streamdelta_media_audit.json"
    if args.check_only:
        report = audit_prepared(output_root, dataset_path)
        write_json_atomic(audit_path, report, pretty=True)
        print(json.dumps(report, ensure_ascii=False))
        return
    if args.jobs < 1:
        parser.error("--jobs must be positive")
    if args.proactive_tail_seconds < 0:
        parser.error("--proactive-tail-seconds must be non-negative")
    if not source_root.is_dir():
        parser.error(f"source root not found: {source_root}")

    output_root.mkdir(parents=True, exist_ok=True)
    ffmpeg = resolve_ffmpeg(args.ffmpeg)
    samples, sources, required_seconds = prepare_samples(
        source_root,
        output_root,
        args.proactive_tail_seconds,
        ffprobe=resolve_ffprobe(ffmpeg),
        probe_workers=min(args.jobs, 32),
    )
    for key, source in sources.items():
        source["required_seconds"] = required_seconds[key]
    manifest = {
        "format": "streamdelta-streamingbench-v1",
        "status": "preparing",
        "source_root": str(source_root),
        "output_root": str(output_root),
        "rows": len(samples),
        "source_videos": len(sources),
        "fps": 1,
        "proactive_tail_seconds": args.proactive_tail_seconds,
        "annotation_sha256": {
            filename: sha256_file(source_root / "StreamingBench" / filename)
            for filename in EXPECTED_CSV_ROWS
        },
        "suite_counts": dict(collections.Counter(
            sample["suite"] for sample in samples
        )),
        "task_counts": dict(collections.Counter(
            sample["task"] for sample in samples
        )),
    }
    write_json_atomic(output_root / "sources.json", sources, pretty=True)
    write_json_atomic(dataset_path, samples)
    write_json_atomic(
        dataset_path.with_suffix(".manifest.json"),
        manifest,
        pretty=True,
    )
    audio_root = output_root / "audio_videos"
    for sample in samples:
        replace_symlink(
            audio_root / f"{sample['id']}.mp4",
            Path(sample["source_video"]),
        )

    frame_counts: Dict[str, int] = {}
    if not args.skip_frames:
        jobs = []
        ordered_sources = sorted(
            sources.items(),
            key=lambda item: (
                int(item[1]["sample_number"]),
                str(item[1]["media_directory"]),
            ),
        )
        for key, source in ordered_sources:
            source_path = Path(source["video"])
            frame_dir = Path(source["frames"])
            duration = required_seconds[key]
            jobs.append((
                key,
                lambda source_path=source_path,
                frame_dir=frame_dir,
                duration=duration: extract_frames(
                    ffmpeg=ffmpeg,
                    source=source_path,
                    frame_dir=frame_dir,
                    duration=duration,
                    force=args.force,
                ),
            ))
        frame_counts = run_frame_jobs(jobs, args.jobs)
    else:
        for key, source in sources.items():
            frame_dir = Path(source["frames"])
            count = completed_frame_count(
                frame_dir,
                Path(source["video"]),
                required_seconds[key],
            )
            if not count:
                raise FileNotFoundError(
                    f"no completed frame extraction for {key}: {frame_dir}"
                )
            frame_counts[key] = count

    for key, source in sources.items():
        source["frame_count"] = frame_counts[key]
    reconcile_sample_times(
        samples,
        frame_counts,
        args.proactive_tail_seconds,
    )

    write_json_atomic(output_root / "sources.json", sources, pretty=True)
    write_json_atomic(dataset_path, samples)
    manifest["status"] = "ready"
    manifest["prepared_source_videos"] = len(frame_counts)
    write_json_atomic(
        dataset_path.with_suffix(".manifest.json"),
        manifest,
        pretty=True,
    )
    report = audit_prepared(output_root, dataset_path)
    write_json_atomic(audit_path, report, pretty=True)
    print(
        f"[prepare] complete: dataset={dataset_path}, rows={len(samples)}, "
        f"source_videos={len(sources)}, audit=passed"
    )


if __name__ == "__main__":
    main()