"""Extract B1 frozen Qwen local crops + recent mixed audio, aligned to Omni .pt.

Uses provided boxes (not speaking labels), raw source media and the existing
dense-timeline reader. No Thinker forward, cross-person attention, or learned
256D projection is performed here. See docs/local_feature_baseline.md.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import tempfile
from pathlib import Path
from typing import Any

import numpy as np
import torch
from tqdm.auto import tqdm

from easycom_head import load_ego4d_manifest_split, normalize_sample_name, resolve_feature_path
from local_turn_features import (
    LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
    load_local_features, pooled_recent_visual, reference_alignment,
)


TESTED_TRANSFORMERS_VERSIONS = ("4.52.3", "4.57.6")


def validate_local_transformers_version(version: str) -> str:
    """Keep the local encoder's tested API separate from the replay KV API."""
    if version not in TESTED_TRANSFORMERS_VERSIONS:
        supported = ", ".join(TESTED_TRANSFORMERS_VERSIONS)
        raise RuntimeError(
            f"Local encoder API is tested with transformers versions {supported}; found {version}. "
            "Use a tested version or validate the encoder APIs before adding support."
        )
    return version


def model_config_digest(config: dict[str, Any]) -> str:
    """Ignore location fields so an explicitly identified checkpoint can move."""
    def stable(value):
        if isinstance(value, dict):
            return {key: stable(item) for key, item in value.items() if key != "_name_or_path"}
        if isinstance(value, (list, tuple)):
            return [stable(item) for item in value]
        return value

    return hashlib.sha256(json.dumps(stable(config), sort_keys=True).encode()).hexdigest()


class FrozenQwenLocalEncoder:
    """Keep only frozen vision+merger and audio+projection on the GPU."""

    def __init__(self, args: argparse.Namespace) -> None:
        import transformers
        from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration

        self.transformers_version = validate_local_transformers_version(transformers.__version__)
        print(f"Local encoders: transformers={self.transformers_version}", flush=True)
        self.device = torch.device(args.device)
        if self.device.type == "cuda" and not torch.cuda.is_available():
            raise RuntimeError(f"CUDA requested but unavailable: {args.device}")
        dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype"
        kwargs = {dtype_key: getattr(torch, args.dtype), "device_map": {"": "cpu"},
                  "output_loading_info": True}
        if args.flash_attention_2:
            kwargs["attn_implementation"] = "flash_attention_2"
        model, info = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.model_id, **kwargs)
        bad = [key for key in info.get("missing_keys", [])
               if key.startswith(("visual.", "audio_tower."))]
        if bad or info.get("mismatched_keys") or info.get("error_msgs"):
            raise RuntimeError(f"Cannot extract frozen features from incompletely loaded weights: {info}")
        model.requires_grad_(False).eval()
        self.visual_dim = int(model.config.vision_config.out_hidden_size)
        self.audio_dim = int(model.config.audio_config.output_dim)
        self.model_config_sha256 = model_config_digest(model.config.to_dict())
        # These modules are never called. Releasing them avoids keeping the
        # multi-billion-parameter language decoder on the extraction GPU.
        model.model = None
        model.lm_head = None
        model.visual.to(self.device)
        model.audio_tower.to(self.device)
        self.model = model
        # Keep PIL/slow image preprocessing consistent with the original B1
        # definition, independently of newer Transformers' fast defaults.
        self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id, use_fast=False)
        self.max_pixels = args.max_pixels
        self.image_batch_size = args.image_batch_size

    @torch.inference_mode()
    def images(self, crops: list[Any]) -> torch.Tensor:
        pieces = []
        for start in range(0, len(crops), self.image_batch_size):
            batch = self.processor.image_processor(
                images=crops[start:start + self.image_batch_size],
                min_pixels=4 * 28 * 28, max_pixels=self.max_pixels, return_tensors="pt",
            )
            grid = batch["image_grid_thw"].to(self.device)
            pixels = batch["pixel_values"].to(self.device, dtype=self.model.visual.dtype)
            tokens = self.model.get_image_features(pixels, grid)
            merge = int(self.model.config.vision_config.spatial_merge_size)
            counts = (grid.prod(-1) // (merge * merge)).tolist()
            if not isinstance(tokens, torch.Tensor) or tokens.shape != (sum(counts), self.visual_dim):
                raise RuntimeError(f"Unexpected image encoder output: {type(tokens).__name__}, "
                                   f"shape={getattr(tokens, 'shape', None)}")
            # Split BEFORE spatial mean: each person has one vote regardless of
            # box size / number of visual tokens.
            pieces.extend(part.float().mean(0).cpu() for part in tokens.split(counts))
        return torch.stack(pieces) if pieces else torch.empty(0, self.visual_dim)

    @torch.inference_mode()
    def audio(self, waveform: np.ndarray) -> torch.Tensor:
        if waveform.ndim != 1 or not len(waveform) or not np.isfinite(waveform).all():
            raise ValueError("Audio must be finite nonempty mono samples")
        # The waveform has ALREADY been cut at the current time. Computing its
        # mel transform here cannot read future samples of the source file.
        batch = self.processor.feature_extractor(
            waveform, sampling_rate=16000, padding="longest",
            return_attention_mask=True, return_tensors="pt",
        )
        features = batch["input_features"].to(self.device, dtype=self.model.audio_tower.dtype)
        mask = batch["attention_mask"].to(self.device)
        tokens = self.model.get_audio_features(features, feature_attention_mask=mask)
        _, lengths = self.model.audio_tower._get_feat_extract_output_lengths(mask.sum(-1))
        if (not isinstance(tokens, torch.Tensor)
                or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens)):
            raise RuntimeError(f"Unexpected valid audio encoder output: {type(tokens).__name__}, "
                               f"shape={getattr(tokens, 'shape', None)}")
        # get_audio_features removes padding and verifies the output length.
        return tokens.float().mean(0).cpu()


def ego4d_clip_boxes(clip: dict[str, Any]) -> dict[int, list[dict[str, Any]]]:
    """AVD tracks indexed by clip-local decoded frame; never use voice labels."""
    result: dict[int, list[dict[str, Any]]] = {}
    for person in clip.get("persons", []):
        if person.get("camera_wearer", False):
            continue
        seen = set()
        for track in person.get("tracking_paths", []):
            for box in track.get("track", []):
                frame = box.get("clip_frame")
                if frame is None:
                    frame = box.get("frame")
                if frame is None:
                    frame = int(box["video_frame"]) - int(clip["video_start_frame"])
                frame = int(frame)
                if frame < 0:
                    raise ValueError("Negative clip-relative Ego4D box frame")
                # Some exports duplicate a person's boundary box in two tracks.
                if frame in seen:
                    continue
                seen.add(frame)
                x, y, w, h = (float(box[k]) for k in ("x", "y", "width", "height"))
                result.setdefault(frame, []).append(
                    {"id": str(person["person_id"]), "xyxy": [x, y, x + w, y + h]}
                )
    return result


def load_ego4d_boxes(paths: list[Path], names: list[str]) -> dict[str, Any]:
    result = {}
    selected = set(names)
    for path in paths:
        print(f"Reading provided tracks: {path}", flush=True)
        with path.open() as handle:
            payload = json.load(handle)
        for video in payload["videos"]:
            for clip in video["clips"]:
                uid = clip["clip_uid"]
                if uid in selected:
                    if uid in result:
                        raise ValueError(f"Duplicate Ego4D annotation UID: {uid}")
                    result[uid] = ego4d_clip_boxes(clip)
        del payload
    missing = set(names) - result.keys()
    if missing:
        raise ValueError(f"No provided tracks for: {sorted(missing)}")
    return result


def build_media_and_boxes(args: argparse.Namespace, name: str, end: float, ego_boxes: dict[str, Any]):
    # Lazy imports allow --help and pure tests without Transformers/decoders.
    from streaming_turn_model_directly import CausalSourceAV, SourceClipSpec, _easycom_directory_specs

    if args.dataset == "easycom":
        specs = _easycom_directory_specs(args.data_root / "Main" / "Video_Compressed" / name)
    else:
        specs = [SourceClipSpec(args.data_root / "v1" / "clips" / f"{name}.mp4")]
    media = CausalSourceAV(specs, output_fps=args.video_fps, required_duration=end,
                           decoder_backend=args.decoder_backend, decoder_threads=args.decoder_threads,
                           decoder_gpu_id=args.decoder_gpu_id)
    if args.dataset == "easycom":
        folder = "Face_Bounding_Boxes" if args.box_kind == "face" else "Head_Bounding_Boxes"
        annotations = []
        for clip in media.clips:
            path = args.data_root / "Main" / folder / name / f"{clip.path.stem}.json"
            with path.open() as handle:
                rows = json.load(handle)
            mapping = easycom_boxes(rows, args.box_kind)
            # Head annotations have a documented extra terminal record; it is
            # not a decoded frame and is never sampled.
            allowed = clip.num_frames + (1 if args.box_kind == "head" else 0)
            if set(mapping) != set(range(allowed)):
                raise ValueError(f"Box/video frame inventory mismatch: {path}")
            annotations.append(mapping)
    else:
        annotations = [ego_boxes[name]]
        if annotations[0] and max(annotations[0]) >= media.clips[0].num_frames:
            raise ValueError(f"Ego4D annotation frame exceeds source clip: {name}")
    return media, annotations


def box_coordinate_offset(clip, coordinate_space: str) -> tuple[int, int]:
    """Map source/encoded boxes into the image actually returned by the reader."""
    if coordinate_space == "encoded":
        return tuple(clip.crop[:2]) if clip.crop is not None else (0, 0)
    if coordinate_space == "camera":
        # Only for annotations already translated into the cropped image.
        return (0, 0)
    raise ValueError(f"Unknown box coordinate space: {coordinate_space}")


def frame_crops(media, annotations, time: float, args: argparse.Namespace):
    clip_index, frame_index, source_pts = media._source_frame_for_time(time)
    clip = media.clips[clip_index]
    image = media.frame_at_or_before(time)
    # EasyCom's supplied boxes refer to the encoded image, including the left
    # banner. The reader has already removed that banner. Subtract this clip's
    # actual crop origin, not a hard-coded offset shared by all sessions.
    offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space)
    crops, records = [], []
    for person in annotations[clip_index].get(frame_index, []):
        box = crop_box(person["xyxy"], *image.size, offset_x=offset_x, offset_y=offset_y,
                       expansion=args.crop_expansion)
        if box is not None:
            crops.append(image.crop(box))
            records.append({"id": person["id"], "box": box,
                            "annotation_box": list(person["xyxy"]),
                            "coordinate_offset": [offset_x, offset_y],
                            "source_clip_index": clip_index,
                            "source_frame_index": frame_index,
                            "source_time": clip.dense_start + source_pts})
    return image, crops, records


def validate_source_timeline(reference: dict[str, Any], media) -> None:
    """Check source order/frames/offsets when the Omni file has direct metadata."""
    saved = (reference.get("direct_media") or {}).get("source_clips")
    if saved is None:
        print("NOTE: reference has no direct-media inventory; verify its dense source timeline manually.")
        return
    if len(saved) < len(media.clips):
        raise ValueError("Local source has more clips than the reference inventory")
    for old, new in zip(saved, media.metadata()["source_clips"]):
        if (Path(old["path"]).name != Path(new["path"]).name
                or old["source_frames"] != new["source_frames"]
                or abs(old["dense_start"] - new["dense_start"]) > 1e-5
                or abs(old["duration"] - new["duration"]) > 1e-5):
            raise ValueError(f"Local/reference source timeline differs at {new['path']}")


def preview_time_bounds(args: argparse.Namespace, available_end: float) -> tuple[float, float]:
    """Preview-only range on the original timeline, never a rebased feature grid."""
    start = float(getattr(args, "preview_start", 0.))
    requested_end = getattr(args, "preview_end", None)
    end = min(available_end, requested_end) if requested_end is not None else available_end
    if start >= end:
        raise ValueError(f"Preview start {start:g}s is outside the available prefix ending at {end:g}s; "
                         "remove/increase --max-steps or choose an earlier --preview-start")
    return start, end


def preview_ticks(media, annotations, args: argparse.Namespace, end: float) -> tuple[list[int], int, int]:
    start, stop = preview_time_bounds(args, end)
    ticks = [tick for tick in causal_frame_ticks(stop, args.video_fps)
             if tick / args.video_fps >= start - 1e-8]
    grid_count = len(ticks)
    if getattr(args, "preview_with_boxes", False):
        eligible = []
        # Scan annotation metadata, not every RGB frame. Use the same source
        # frame mapping, coordinate correction and crop validity as extraction.
        for tick in ticks:
            clip_index, frame_index, _ = media._source_frame_for_time(tick / args.video_fps)
            clip = media.clips[clip_index]
            width, height = clip.crop[2:] if clip.crop is not None else (clip.width, clip.height)
            offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space)
            if any(crop_box(person["xyxy"], width, height, offset_x=offset_x, offset_y=offset_y,
                            expansion=args.crop_expansion) is not None
                   for person in annotations[clip_index].get(frame_index, [])):
                eligible.append(tick)
        ticks = eligible
    eligible_count = len(ticks)
    count = min(getattr(args, "preview_count", 4), eligible_count)
    indices = np.linspace(0, eligible_count - 1, count, dtype=int) if count else []
    return [ticks[index] for index in indices], grid_count, eligible_count


def preview_boxes(media, annotations, args: argparse.Namespace, name: str, end: float) -> None:
    from PIL import ImageDraw

    folder = args.output_root / f"{name}_preview"
    folder.mkdir(parents=True, exist_ok=True)
    start, stop = preview_time_bounds(args, end)
    ticks, grid_count, eligible_count = preview_ticks(media, annotations, args, end)
    print(f"Preview [{start:g},{stop:g})s: {grid_count} grid frames, {eligible_count} eligible, "
          f"saving {len(ticks)} images", flush=True)
    if not ticks:
        print("No eligible preview frames. Try a longer range or another clip. "
              "If people are visible but no valid boxes exist, check annotation/frame alignment.", flush=True)
    audit = []
    for tick in ticks:
        time = tick / args.video_fps
        image, _, records = frame_crops(media, annotations, time, args)
        clip_index, frame_index, _ = media._source_frame_for_time(time)
        clip = media.clips[clip_index]
        offset = box_coordinate_offset(clip, args.box_coordinate_space)
        draw = ImageDraw.Draw(image)
        for record in records:
            draw.rectangle(record["box"], outline="lime", width=4)
            draw.text(record["box"][:2], record["id"], fill="red")
        caption = f"t={time:.2f}s | source frame={frame_index} | people={len(records)}"
        draw.rectangle(draw.textbbox((8, 8), caption), fill="black")
        draw.text((8, 8), caption, fill="white")
        filename = f"frame_{tick:06d}.jpg"
        image.save(folder / filename)
        audit.append({
            "image": filename, "time": time, "source_clip": str(clip.path),
            "source_frame_index": frame_index, "source_size": [clip.width, clip.height],
            "crop": list(clip.crop) if clip.crop is not None else None,
            "box_coordinate_space": args.box_coordinate_space,
            "coordinate_offset": list(offset), "boxes": records,
        })
        print(f"preview t={time:g}s: {clip.width}x{clip.height} -> {image.width}x{image.height}; "
              f"boxes={args.box_coordinate_space}, subtract offset={offset}, people={len(records)}")
    (folder / "boxes.json").write_text(json.dumps(audit, indent=2) + "\n")
    (folder / "preview.json").write_text(json.dumps({
        "start": start, "end": stop, "video_fps": args.video_fps,
        "with_boxes_only": getattr(args, "preview_with_boxes", False),
        "requested_count": getattr(args, "preview_count", 4),
        "grid_frame_count": grid_count, "eligible_frame_count": eligible_count,
        "saved_count": len(audit), "saved_with_boxes": sum(bool(row["boxes"]) for row in audit),
        "images": [row["image"] for row in audit],
    }, indent=2) + "\n")
    print(f"Box previews: {folder} (inspect before extracting all sessions)")
    print("boxes.json lists this run's images; older images in the directory are not removed.")


def extract_local(media, annotations, encoder, times: torch.Tensor, args: argparse.Namespace) -> dict[str, Any]:
    # Canonical grid, not float32-serialized rounding near a tick boundary.
    ends = (torch.arange(len(times), dtype=torch.float64) + 1) / 5
    ticks = list(causal_frame_ticks(float(ends[-1]), args.video_fps))
    frame_times = torch.tensor(ticks, dtype=torch.float64) / args.video_fps
    frame_vectors = torch.zeros(len(ticks), encoder.visual_dim)
    frame_valid = torch.zeros(len(ticks), dtype=torch.bool)
    people, records, offsets = [], [], [0]
    for row, tick in enumerate(tqdm(ticks, desc="local vision", unit="frame", disable=args.no_progress)):
        _, crops, observations = frame_crops(media, annotations, tick / args.video_fps, args)
        vectors = encoder.images(crops)
        if vectors.shape != (len(crops), encoder.visual_dim) or not torch.isfinite(vectors).all():
            raise RuntimeError("Invalid person features")
        if len(vectors):
            stored = vectors.to(torch.bfloat16)
            frame_vectors[row] = stored.float().mean(0)
            frame_valid[row] = True
            people.append(stored)
        records.extend(observations)
        offsets.append(offsets[-1] + len(vectors))
    visual = torch.zeros(len(times), encoder.visual_dim)
    audio = torch.zeros(len(times), encoder.audio_dim)
    visual_valid = torch.zeros(len(times), dtype=torch.bool)
    audio_valid = torch.zeros(len(times), dtype=torch.bool)
    for row, end in enumerate(tqdm(ends.tolist(), desc="local audio/pool", unit="state", disable=args.no_progress)):
        visual[row], present = pooled_recent_visual(
            frame_times, frame_vectors, frame_valid, end, args.visual_window_seconds,
        )
        visual_valid[row] = present
        waveform = media.audio_slice(max(0., end - args.audio_window_seconds), end)
        audio[row] = encoder.audio(waveform)
        audio_valid[row] = True  # Silence is an observation, not missing audio.
    if not torch.isfinite(audio).all() or not torch.isfinite(visual).all():
        raise RuntimeError("Non-finite local encoder outputs")
    return {
        "visual": visual.to(torch.bfloat16), "audio": audio.to(torch.bfloat16),
        "visual_valid": visual_valid, "audio_valid": audio_valid,
        "persons": {"features": torch.cat(people) if people else torch.empty(0, encoder.visual_dim, dtype=torch.bfloat16),
                    "frame_times": frame_times, "frame_offsets": torch.tensor(offsets, dtype=torch.long),
                    "records": records},
    }


def select_samples(args: argparse.Namespace) -> list[str]:
    if args.manifest is not None:
        names = load_ego4d_manifest_split(args.manifest, args.split)
        # One-based inclusive range in the selected split, then an optional limit.
        names = names[args.start_index - 1:args.end_index]
        if args.limit is not None:
            names = names[:args.limit]
    else:
        names = [normalize_sample_name(value, args.dataset) for value in args.samples]
    if not names:
        raise ValueError("No samples selected")
    if len(set(names)) != len(names):
        raise ValueError("Duplicate samples")
    num_shards = getattr(args, "num_shards", 1)
    shard_index = getattr(args, "shard_index", 0)
    names = names[shard_index::num_shards]
    return names


def requested_local_config(args: argparse.Namespace) -> dict[str, Any]:
    """Settings that must match before reusing an existing complete sidecar."""
    return {
        "encoder_id": args.encoder_id or args.model_id,
        "dtype": args.dtype,
        "attention": "flash_attention_2" if args.flash_attention_2 else "default",
        "visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds,
        "audio_window_seconds": args.audio_window_seconds,
        "visual_pooling": "spatial-mean_then_equal-person-mean_then_equal-valid-frame-mean",
        "audio_pooling": "valid-output-token-mean", "audio_source": "video-mixed-mono-16khz",
        "box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks",
        "box_coordinate_space": args.box_coordinate_space, "crop_expansion": args.crop_expansion,
        "max_pixels": args.max_pixels, "uses_speech_labels": False,
        "causal_policy": "half-open-trailing-window", "box_protocol": "provided-boxes-not-online-detector",
    }


def validate_reusable_sidecar(args, name, reference, version, expected_contract=None):
    path = args.output_root / f"{name}.pt"
    _, contract = load_local_features(path, reference, sample_id=name, dataset=args.dataset,
                                      modalities="av", expected_contract=expected_contract)
    config = contract["config"]
    expected = requested_local_config(args)
    if version is not None:
        expected["transformers"] = version
    mismatches = [key for key, value in expected.items() if config.get(key) != value]
    if mismatches or not config.get("model_config_sha256") or not config.get("transformers"):
        raise ValueError(f"Cannot reuse {path}: incompatible/incomplete encoder settings {mismatches}; "
                         "use a new output root or --overwrite")
    return contract


def parse_args(argv=None):
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
    p.add_argument("--data-root", required=True, type=Path)
    p.add_argument("--feature-root", required=True, type=Path,
                   help="Flat .pt directory or Ego4D parent containing train/ and test/")
    p.add_argument("--target-root", required=True, type=Path)
    p.add_argument("--output-root", required=True, type=Path)
    selection = p.add_mutually_exclusive_group(required=True)
    selection.add_argument("--samples", nargs="+", help="EasyCom session numbers or Ego4D clip UIDs")
    selection.add_argument("--manifest", type=Path, help="Ego4D EgoSpeak TSV; never creates a new split")
    p.add_argument("--split", choices=("all", "train", "test"), default="all")
    p.add_argument("--start-index", type=int, default=1, help="One-based index AFTER split filtering")
    p.add_argument("--end-index", type=int, help="Inclusive index AFTER split filtering")
    p.add_argument("--limit", type=int, help="First N clips in the selected manifest range")
    p.add_argument("--num-shards", type=int, default=1, help="Disjoint round-robin shards after all sample selection")
    p.add_argument("--shard-index", type=int, default=0, help="Zero-based shard index")
    p.add_argument("--ego4d-annotations", type=Path, nargs="+")
    p.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
    p.add_argument("--encoder-id", help="Stable checkpoint identity for relocated copies; default: model-id")
    p.add_argument("--device", default="cuda:0")
    p.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16")
    p.add_argument("--flash-attention-2", action="store_true")
    p.add_argument("--video-fps", type=float, default=2.)
    p.add_argument("--audio-window-seconds", type=float, default=1.)
    p.add_argument("--visual-window-seconds", type=float, default=1.)
    p.add_argument("--box-kind", choices=("face", "head"), default="face")
    p.add_argument(
        "--box-coordinate-space", choices=("camera", "encoded"), default=None,
        help=("Box coordinates before or after video cropping. Default: encoded for EasyCom "
              "(subtract the left banner), camera for uncropped Ego4D. Use camera only for "
              "annotations already translated to the cropped image."),
    )
    p.add_argument("--crop-expansion", type=float, default=1.)
    p.add_argument("--max-pixels", type=int, default=224 * 224)
    p.add_argument("--image-batch-size", type=int, default=8)
    p.add_argument("--decoder-backend", choices=("decord", "torchcodec", "pynvcodec"), default="decord")
    p.add_argument("--decoder-threads", type=int, default=1)
    p.add_argument("--decoder-gpu-id", type=int, default=0, help="Visible-device index, not physical GPU number")
    p.add_argument("--max-steps", type=int, help="Smoke test only; partial sidecars are rejected by training")
    p.add_argument("--preview-only", action="store_true", help="Save box overlays without loading the model")
    p.add_argument("--preview-count", type=int, default=4, help="Maximum images per clip (preview only)")
    p.add_argument("--preview-start", type=float, default=0., help="Start time in seconds on the clip/dense timeline")
    p.add_argument("--preview-end", type=float, help="Exclusive end time; default: entire available prefix")
    p.add_argument("--preview-with-boxes", action="store_true",
                   help="Sample only grid frames with valid provided crops; preview selection only, not feature filtering")
    existing = p.add_mutually_exclusive_group()
    existing.add_argument("--overwrite", action="store_true")
    existing.add_argument("--reuse-existing", action="store_true", help="Validate and skip complete matching sidecars")
    p.add_argument("--dry-run", action="store_true", help="Validate references/targets and list work; no model, media decode or writes")
    p.add_argument("--no-progress", action="store_true")
    args = p.parse_args(argv)
    if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards:
        p.error("Require --num-shards > 0 and 0 <= --shard-index < --num-shards")
    if args.manifest is not None and args.dataset != "ego4d":
        p.error("--manifest is only supported with --dataset ego4d")
    if args.manifest is None and (args.split != "all" or args.start_index != 1
                                  or args.end_index is not None or args.limit is not None):
        p.error("--split/--start-index/--end-index/--limit require --manifest")
    if args.start_index < 1 or (args.end_index is not None and args.end_index < args.start_index):
        p.error("Invalid manifest index range")
    if args.limit is not None and args.limit <= 0:
        p.error("--limit must be positive")
    if args.reuse_existing and (args.max_steps is not None or args.preview_only):
        p.error("--reuse-existing is for full extraction, not smoke tests or previews")
    if args.dry_run and args.preview_only:
        p.error("--dry-run and --preview-only are separate checks")
    if args.preview_count <= 0:
        p.error("--preview-count must be positive")
    if not np.isfinite(args.preview_start) or args.preview_start < 0:
        p.error("--preview-start must be finite and nonnegative")
    if args.preview_end is not None and (not np.isfinite(args.preview_end) or args.preview_end <= args.preview_start):
        p.error("--preview-end must be finite and greater than --preview-start")
    if not args.preview_only and (args.preview_count != 4 or args.preview_start != 0.
                                   or args.preview_end is not None or args.preview_with_boxes):
        p.error("Preview selection options require --preview-only")
    if args.preview_only and args.max_steps is not None:
        if args.preview_start >= args.max_steps / 5 or (args.preview_end is not None
                                                       and args.preview_end > args.max_steps / 5):
            p.error("Preview range exceeds --max-steps; remove --max-steps to inspect later frames")
    if args.box_coordinate_space is None:
        args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera"
    for name in ("video_fps", "audio_window_seconds", "visual_window_seconds", "crop_expansion"):
        if not np.isfinite(getattr(args, name)) or getattr(args, name) <= 0:
            p.error(f"--{name.replace('_', '-')} must be finite and positive")
    if args.audio_window_seconds > 30:
        p.error("--audio-window-seconds must not exceed the audio processor's 30s limit")
    if args.max_pixels < 4 * 28 * 28 or args.image_batch_size <= 0:
        p.error("max-pixels must be >=3136; image-batch-size must be positive")
    if args.max_steps is not None and args.max_steps <= 0:
        p.error("--max-steps must be positive")
    protected_roots = [args.feature_root]
    if args.dataset == "ego4d":
        protected_roots += [args.feature_root / split for split in ("train", "test")]
    if args.output_root.resolve() in [root.resolve() for root in protected_roots]:
        p.error("Local outputs must not overwrite original Omni features")
    return args


def main(argv=None) -> None:
    args = parse_args(argv)
    version = None
    if not args.preview_only and not args.dry_run:
        import transformers

        # Fail before indexing/decoding source media, not after expensive I/O.
        version = validate_local_transformers_version(transformers.__version__)
    names = select_samples(args)
    if not names:
        print(f"Shard {args.shard_index}/{args.num_shards}: no assigned clips")
        return
    print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned clips", flush=True)
    reference_paths = {name: resolve_feature_path(args.feature_root, name, dataset=args.dataset) for name in names}
    paths = []
    if args.dataset == "ego4d":
        paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
                                          for split in ("train", "val")]
        for path in paths:
            if not path.is_file():
                raise FileNotFoundError(path)
    for name in names:
        required = [reference_paths[name], args.target_root / f"{name}.npy"]
        if args.dataset == "ego4d":
            required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
        for path in required:
            if not path.is_file():
                raise FileNotFoundError(path)
        if ((args.output_root / f"{name}.pt").exists() and not args.overwrite
                and not args.preview_only and not args.reuse_existing):
            raise FileExistsError(f"{name}.pt exists; use --reuse-existing, a new output root or --overwrite")
    pending, contract = [], None
    # Stream through references on CPU; never keep all global states in memory.
    for name in tqdm(names, desc="local preflight", unit="clip", disable=args.no_progress or args.dry_run):
        reference = torch.load(reference_paths[name], map_location="cpu", weights_only=False)
        times, _ = reference_alignment(reference)
        target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
        if target.shape != (len(times), 3):
            raise ValueError(f"Reference/target mismatch: {name}, {target.shape} vs {len(times)} states")
        if args.reuse_existing and (args.output_root / f"{name}.pt").is_file():
            contract = validate_reusable_sidecar(args, name, reference, version, contract)
            print(f"{name}: reuse ({len(times)} states)", flush=True)
        else:
            pending.append(name)
            if args.dry_run:
                steps = min(len(times), args.max_steps or len(times))
                print(f"{name}: extract ({steps} states, {steps / 5:g}s)", flush=True)
        del reference, target
    print(f"Selected {len(names)} clips; reuse={len(names) - len(pending)}, pending={len(pending)}", flush=True)
    if args.dry_run:
        print("Dry run only: annotation contents, pixel alignment and model/runtime are not checked.")
        return
    if not pending:
        return
    ego_boxes = load_ego4d_boxes(paths, pending) if args.dataset == "ego4d" else {}
    args.output_root.mkdir(parents=True, exist_ok=True)
    encoder = None
    for position, name in enumerate(pending, 1):
        print(f"[{position}/{len(pending)}] {name}", flush=True)
        path = reference_paths[name]
        reference = torch.load(path, map_location="cpu", weights_only=False)
        times, indices = reference_alignment(reference)
        target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
        if target.shape != (len(times), 3):
            raise ValueError(f"Reference/target mismatch: {name}, {target.shape} vs {len(times)} states")
        full_steps = len(times)
        if args.max_steps is not None:
            times, indices = times[:args.max_steps], indices[:args.max_steps]
        media_end = len(times) / 5
        if args.preview_only:
            _, media_end = preview_time_bounds(args, media_end)
        media, annotations = build_media_and_boxes(args, name, media_end, ego_boxes)
        validate_source_timeline(reference, media)
        if (args.dataset == "easycom" and args.box_coordinate_space == "camera"
                and any(clip.crop is not None and clip.crop[0] for clip in media.clips)):
            print("NOTE: camera boxes receive NO banner-offset correction. For original EasyCom "
                  "annotations use --box-coordinate-space encoded; otherwise boxes shift right.", flush=True)
        if args.preview_only:
            preview_boxes(media, annotations, args, name, len(times) / 5)
            ego_boxes.pop(name, None)
            del media, annotations, reference
            continue
        if encoder is None:
            encoder = FrozenQwenLocalEncoder(args)
        config = {**requested_local_config(args),
                  "model_config_sha256": encoder.model_config_sha256,
                  "transformers": encoder.transformers_version,
                  "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim}
        new_contract = {"format": LOCAL_FORMAT, "modalities": "av", "config": config}
        if contract is not None and contract != new_contract:
            raise ValueError("Current encoder differs from reused sidecars; use a new output root")
        contract = new_contract
        print(f"{name}: {len(times)} states; native visual={encoder.visual_dim}D audio={encoder.audio_dim}D", flush=True)
        output = extract_local(media, annotations, encoder, times, args)
        output.update({
            "format": LOCAL_FORMAT, "dataset": args.dataset, "sample_id": name,
            "times": times, "label_indices": indices, "partial": len(times) != full_steps,
            "reference_features": str(path.resolve()), "direct_media": media.metadata(),
            "config": config,
        })
        destination = args.output_root / f"{name}.pt"
        # Atomic replacement: a failed extraction never looks like a ready file.
        with tempfile.NamedTemporaryFile(dir=args.output_root, prefix=f".{name}.", suffix=".tmp", delete=False) as handle:
            temporary = Path(handle.name)
        try:
            torch.save(output, temporary)
            os.replace(temporary, destination)
        finally:
            temporary.unlink(missing_ok=True)
        print(f"Saved {destination}; V-valid={output['visual_valid'].float().mean():.1%}, "
              f"A-valid={output['audio_valid'].float().mean():.1%}, partial={output['partial']}", flush=True)
        ego_boxes.pop(name, None)
        del media, annotations, reference, output


if __name__ == "__main__":
    main()