#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0

from __future__ import annotations

import errno
import json
import os
import shutil
import threading
import weakref
import zipfile
from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from pathlib import Path
from typing import Any
from zipfile import ZipFile

import numpy as np
import pytest
from filelock import FileLock

from lerobot.streaming import sidecar_utils
from lerobot.streaming.manifest import EpisodeVideoManifest, VideoFileRecord
from lerobot.streaming.mp4 import Mp4Index
from lerobot.streaming.sidecar import (
    SidecarLockTimeoutError,
    SidecarSpec,
    ensure_mp4_sidecar,
    sidecar_cache_path,
)


@pytest.fixture(autouse=True)
def isolated_index_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HF_LEROBOT_HOME", str(tmp_path / "cache-home"))


def _record(path: str = "videos/camera/chunk-000/file-000.mp4", size: int = 128) -> VideoFileRecord:
    arrays = np.array([0], dtype=np.int64)
    index = Mp4Index(
        file_path=path,
        file_size=size,
        ftyp=b"",
        moov_offset=0,
        mdat_offset=0,
        mdat_payload_offset=0,
        mdat_payload_size=size,
        faststart=True,
        codec="avc1",
        timescale=1,
        duration=1,
        track_id=1,
        width=1,
        height=1,
        stsd_body=b"",
        sample_pts=np.array([0.0]),
        sample_durations=arrays,
        sample_composition_offsets=arrays,
        sample_sizes=arrays,
        sample_offsets=arrays,
        sync_samples=arrays,
    )
    return VideoFileRecord(path, size, index)


def _spec(revision: str = "rev-a", size: int = 128) -> SidecarSpec:
    return SidecarSpec(
        repo_id="owner/dataset",
        revision=revision,
        data_root="hf://datasets/owner/dataset",
        source_files=(("videos/camera/chunk-000/file-000.mp4", size),),
    )


def _write_valid(path: Path, spec: SidecarSpec) -> None:
    EpisodeVideoManifest.save_file_sidecar(path, [_record(size=spec.source_files[0][1])], spec=spec)


def test_sidecar_cache_path_is_revision_keyed(tmp_path: Path) -> None:
    first = sidecar_cache_path(tmp_path, _spec("rev-a"))
    second = sidecar_cache_path(tmp_path, _spec("rev-b"))

    assert first != second
    assert first.parent == second.parent


@pytest.mark.parametrize("root", ["hf://datasets/owner/dataset@main", "hf://buckets/owner/data", "/local"])
def test_mutable_or_local_source_still_requires_revision_match(root: str) -> None:
    expected = replace(_spec(), data_root=root)
    assert not expected.matches(replace(expected, revision="different"))


def test_existing_alias_cache_is_reused_without_rebuild(tmp_path: Path) -> None:
    expected = replace(_spec("v3.0"), data_root=f"hf://datasets/owner/dataset@{'a' * 40}")
    original_path = sidecar_cache_path(tmp_path, expected)
    _write_valid(original_path, expected)

    def unexpected_build(path: Path, spec: SidecarSpec) -> None:
        pytest.fail("A valid existing local index must not be downloaded or rebuilt")

    def unexpected_download(path: Path, spec: SidecarSpec) -> bool:
        pytest.fail("A valid existing local index must not be downloaded or rebuilt")

    assert (
        ensure_mp4_sidecar(expected, tmp_path, build=unexpected_build, download=unexpected_download)
        == original_path
    )


def test_source_generation_round_trip_and_replacement(tmp_path: Path) -> None:
    original = _spec()
    video = original.source_files[0][0]
    first = replace(original, source_fingerprints=((video, "hash-one"),))
    second = replace(original, source_fingerprints=((video, "hash-two"),))
    built = []

    def build(path: Path, spec: SidecarSpec) -> None:
        built.append(spec)
        _write_valid(path, spec)

    first_path = ensure_mp4_sidecar(first, tmp_path, build=build)
    assert ensure_mp4_sidecar(first, tmp_path, build=build) == first_path
    second_path = ensure_mp4_sidecar(second, tmp_path, build=build)
    assert built == [first, second]
    assert first_path != second_path
    assert SidecarSpec.from_dict(second.to_dict()) == second
    assert not EpisodeVideoManifest.validate_file_sidecar(first_path, second)
    assert EpisodeVideoManifest.validate_file_sidecar(second_path, second)
    assert EpisodeVideoManifest.validate_file_sidecar(first_path, first)


def test_ensure_reuses_valid_local_sidecar(tmp_path: Path) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    _write_valid(path, spec)
    build_calls = 0

    def build(_path: Path, _spec: SidecarSpec) -> None:
        nonlocal build_calls
        build_calls += 1

    resolved = ensure_mp4_sidecar(spec, tmp_path, build=build)

    assert resolved == path
    assert build_calls == 0


def test_ensure_rebuilds_sidecar_without_composition_timing(tmp_path: Path) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    path.parent.mkdir(parents=True, exist_ok=True)
    old_spec = SidecarSpec(spec.repo_id, spec.revision, spec.data_root, spec.source_files, schema_version=2)
    # Schema 2 recorded decode timestamps as presentation timestamps and cannot
    # safely be reused even when its source identity and file sizes match.
    np.savez_compressed(
        path, manifest_json=json.dumps({"version": 2, "sidecar": old_spec.to_dict()}).encode()
    )
    assert path != sidecar_cache_path(tmp_path, old_spec)
    assert not EpisodeVideoManifest.validate_file_sidecar(path, spec)
    with pytest.raises(ValueError, match="Unsupported MP4 sidecar schema"):
        EpisodeVideoManifest.load_file_sidecar(path)
    build_calls = []

    def build(target: Path, target_spec: SidecarSpec) -> None:
        build_calls.append(target_spec)
        _write_valid(target, target_spec)

    assert ensure_mp4_sidecar(spec, tmp_path, build=build) == path
    assert build_calls == [spec]
    assert EpisodeVideoManifest.validate_file_sidecar(path, spec)


def test_ensure_downloads_valid_published_sidecar(tmp_path: Path) -> None:
    spec = _spec()
    published = tmp_path / "published.npz"
    _write_valid(published, spec)
    build_calls = 0

    def download(path: Path, _spec: SidecarSpec) -> bool:
        shutil.copyfile(published, path)
        return True

    def build(_path: Path, _spec: SidecarSpec) -> None:
        nonlocal build_calls
        build_calls += 1

    resolved = ensure_mp4_sidecar(spec, tmp_path / "cache", build=build, download=download)

    assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec)
    assert build_calls == 0


@pytest.mark.parametrize("invalid_kind", ["corrupt", "stale"])
def test_ensure_rebuilds_invalid_local_sidecar(tmp_path: Path, invalid_kind: str) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    path.parent.mkdir(parents=True, exist_ok=True)
    if invalid_kind == "corrupt":
        path.write_bytes(b"not-an-npz")
    else:
        _write_valid(path, _spec(revision="other-revision"))
    build_calls = 0

    def build(target: Path, target_spec: SidecarSpec) -> None:
        nonlocal build_calls
        build_calls += 1
        _write_valid(target, target_spec)

    resolved = ensure_mp4_sidecar(spec, tmp_path, build=build)

    assert EpisodeVideoManifest.validate_file_sidecar(resolved, spec)
    assert build_calls == 1


def test_concurrent_ensure_builds_once(tmp_path: Path) -> None:
    spec = _spec()
    start = threading.Barrier(2)
    build_calls = 0
    build_lock = threading.Lock()

    def build(path: Path, target_spec: SidecarSpec) -> None:
        nonlocal build_calls
        with build_lock:
            build_calls += 1
        _write_valid(path, target_spec)

    def ensure() -> Path:
        start.wait()
        return ensure_mp4_sidecar(spec, tmp_path, build=build)

    with ThreadPoolExecutor(max_workers=2) as pool:
        paths = list(pool.map(lambda _: ensure(), range(2)))

    assert paths[0] == paths[1]
    assert build_calls == 1


def test_failed_build_does_not_replace_existing_file(tmp_path: Path) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(b"old-corrupt-file")

    def build(target: Path, _spec: SidecarSpec) -> None:
        target.write_bytes(b"partial")
        raise RuntimeError("build failed")

    with pytest.raises(RuntimeError, match="build failed"):
        ensure_mp4_sidecar(spec, tmp_path, build=build)

    assert path.read_bytes() == b"old-corrupt-file"
    assert not list(path.parent.glob(f".{path.name}.*.tmp.npz"))


def test_lock_timeout_is_actionable(tmp_path: Path) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    lock_path = path.with_suffix(f"{path.suffix}.lock")
    lock_path.parent.mkdir(parents=True, exist_ok=True)

    with (
        FileLock(lock_path),
        pytest.raises(SidecarLockTimeoutError, match="Timed out waiting"),
    ):
        ensure_mp4_sidecar(spec, tmp_path, build=_write_valid, lock_timeout_s=0.01)


def test_sidecar_arrays_are_read_only_file_backed(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    first = EpisodeVideoManifest.load_file_sidecar(path)
    second = EpisodeVideoManifest.load_file_sidecar(path)
    for records in (first, second):
        array = next(iter(records.values())).mp4.sample_pts
        assert not array.flags.writeable
        assert isinstance(array.base, np.memmap)
        np.testing.assert_array_equal(array, [0.0])


def test_warm_validation_does_not_decompress_arrays(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    assert EpisodeVideoManifest.validate_file_sidecar(path, _spec())

    def unexpected_load(*args: object, **kwargs: object) -> None:
        pytest.fail("A warm index must not reopen the compressed sidecar")

    monkeypatch.setattr(np, "load", unexpected_load)
    assert EpisodeVideoManifest.validate_file_sidecar(path, _spec())
    records = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(records.values())).mp4.sample_pts, [0.0])


def test_sidecar_replacement_preserves_live_arrays(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    old = EpisodeVideoManifest.load_file_sidecar(path)
    replacement = _record()
    replacement.mp4.sample_pts[0] = 42.0
    EpisodeVideoManifest.save_file_sidecar(path, [replacement], spec=_spec())
    new = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(old.values())).mp4.sample_pts, [0.0])
    np.testing.assert_array_equal(next(iter(new.values())).mp4.sample_pts, [42.0])


@pytest.mark.parametrize("warm", [False, True])
def test_sidecar_generation_uses_open_file_not_stale_path_attributes(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, warm: bool
) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    stale_stat = path.stat()
    old = EpisodeVideoManifest.load_file_sidecar(path) if warm else None
    replacement = _record()
    replacement.mp4.sample_pts[0] = 42.0
    EpisodeVideoManifest.save_file_sidecar(path, [replacement], spec=_spec())
    original_stat = Path.stat

    def stale_path_stat(self: Path, *args: object, **kwargs: object) -> os.stat_result:
        # Model a shared filesystem's stale pathname attributes, not stale file data.
        return stale_stat if self == path else original_stat(self, *args, **kwargs)

    monkeypatch.setattr(Path, "stat", stale_path_stat)
    new = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(new.values())).mp4.sample_pts, [42.0])
    if old is not None:
        np.testing.assert_array_equal(next(iter(old.values())).mp4.sample_pts, [0.0])


def test_sidecar_payload_does_not_reuse_stale_path_generation(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec("old"))
    stale_stat = path.stat()
    sidecar_utils.mapped_sidecar(path)
    _write_valid(path, _spec("new"))
    original_stat = Path.stat

    def stale_path_stat(self: Path, *args: object, **kwargs: object) -> os.stat_result:
        return stale_stat if self == path else original_stat(self, *args, **kwargs)

    monkeypatch.setattr(Path, "stat", stale_path_stat)
    assert sidecar_utils.sidecar_payload(path)["sidecar"]["revision"] == "new"


def test_rewrite_with_repeated_stat_identity_is_not_served_stale(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Coarse timestamps and inode reuse can give a quick same-size rewrite the old stat identity."""
    path = tmp_path / "index.npz"
    _write_valid(path, _spec("old"))
    old_identity = os.stat(path)
    sidecar_utils.mapped_sidecar(path)
    _write_valid(path, _spec("new"))
    original_fstat = os.fstat

    def repeated_fstat(fd: int) -> os.stat_result:
        stat = original_fstat(fd)
        return old_identity if stat.st_ino == os.stat(path).st_ino else stat

    monkeypatch.setattr(sidecar_utils.os, "fstat", repeated_fstat)
    assert sidecar_utils.sidecar_payload(path)["sidecar"]["revision"] == "new"
    assert sidecar_utils.mapped_sidecar(path)[1]["sidecar"]["revision"] == "new"


def test_source_change_during_conversion_is_not_published(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    validate = sidecar_utils._validate_arrays

    def change_after_array_read(arrays: dict[str, np.ndarray]) -> None:
        validate(arrays)
        # Mutate the real, still-open source after its arrays have been read.
        with path.open("ab") as changed:
            changed.write(b"concurrent mutation")

    monkeypatch.setattr(sidecar_utils, "_validate_arrays", change_after_array_read)
    with pytest.raises(OSError, match="changed during index conversion"):
        sidecar_utils.mapped_sidecar(path)
    cache = tmp_path / "cache-home" / "streaming-indexes"
    assert not list(cache.glob("*.bin"))
    assert not list(cache.glob("*.index.tmp"))


def test_sidecar_load_projects_source_files(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record("a.mp4"), _record("b.mp4")], spec=_spec())
    records = EpisodeVideoManifest.load_file_sidecar(path, file_paths=["b.mp4"])
    assert list(records) == ["b.mp4"]
    np.testing.assert_array_equal(records["b.mp4"].mp4.sample_pts, [0.0])


def test_concurrent_index_conversion_runs_once(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    original = np.load
    calls = []

    def counted_load(*args: object, **kwargs: object) -> object:
        calls.append(1)
        return original(*args, **kwargs)

    monkeypatch.setattr(np, "load", counted_load)
    with ThreadPoolExecutor(max_workers=4) as pool:
        records = list(pool.map(EpisodeVideoManifest.load_file_sidecar, [path] * 4))
    assert len(calls) == 1
    for result in records:
        np.testing.assert_array_equal(next(iter(result.values())).mp4.sample_pts, [0.0])


def test_invalid_derived_index_is_recreated(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    cache_path, _ = sidecar_utils.mapped_sidecar(path)
    cache_path.write_bytes(b"interrupted")
    records = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(records.values())).mp4.sample_pts, [0.0])


@pytest.mark.parametrize("corruption", ["non-object", "fractional-count"])
def test_invalid_derived_metadata_is_recreated(tmp_path: Path, corruption: str) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    cache_path, payload = sidecar_utils.mapped_sidecar(path)
    raw = cache_path.read_bytes()
    old_length = int.from_bytes(raw[-16:-8], "little")
    if corruption == "non-object":
        metadata = b"[]"
    else:
        payload["files"][0]["arrays"]["sample_pts"][1] = 0.5
        metadata = json.dumps(payload).encode()
    cache_path.write_bytes(
        raw[: -16 - old_length] + metadata + len(metadata).to_bytes(8, "little") + raw[-8:]
    )
    records = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(records.values())).mp4.sample_pts, [0.0])


def test_temporary_validation_does_not_prepare_index(tmp_path: Path) -> None:
    path = tmp_path / "temporary.npz"
    _write_valid(path, _spec())
    assert EpisodeVideoManifest.validate_file_sidecar(path, _spec(), prepare_cache=False)
    assert not list((tmp_path / "cache-home").glob("**/*.bin"))


def test_wrong_revision_does_not_prepare_index(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    assert not EpisodeVideoManifest.validate_file_sidecar(path, _spec("wrong"))
    assert not list((tmp_path / "cache-home").glob("**/*.bin"))


def test_missing_array_is_rejected_without_publishing_index(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    with np.load(path, allow_pickle=False) as data:
        arrays = {key: data[key] for key in data.files if key != "0/sample_offsets"}
    np.savez_compressed(path, **arrays)
    assert not EpisodeVideoManifest.validate_file_sidecar(path, _spec())
    assert not list((tmp_path / "cache-home").glob("**/*.bin"))
    assert not list((tmp_path / "cache-home").glob("**/*.tmp"))


def test_read_only_source_uses_writable_local_index_cache(tmp_path: Path) -> None:
    source_dir = tmp_path / "source"
    path = source_dir / "index.npz"
    _write_valid(path, _spec())
    path.chmod(0o444)
    source_dir.chmod(0o555)
    try:
        records = EpisodeVideoManifest.load_file_sidecar(path)
        assert list(source_dir.iterdir()) == [path]
        np.testing.assert_array_equal(next(iter(records.values())).mp4.sample_pts, [0.0])
    finally:
        source_dir.chmod(0o755)


def test_full_index_cache_does_not_rebuild_valid_source(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    _write_valid(path, spec)

    def full_disk(*args: object, **kwargs: object) -> None:
        raise OSError(errno.ENOSPC, "No space left on device")

    def unexpected_build(*args: object) -> None:
        pytest.fail("A full derived cache is not an invalid source sidecar")

    monkeypatch.setattr(sidecar_utils.tempfile, "NamedTemporaryFile", full_disk)
    with pytest.raises(OSError, match="free space in HF_LEROBOT_HOME"):
        ensure_mp4_sidecar(spec, tmp_path, build=unexpected_build)


@pytest.mark.parametrize("prepare_cache", [False, True])
def test_sidecar_array_reads_do_not_scan_npz_member_lists(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, prepare_cache: bool
) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    original = np.lib.npyio.NpzFile.__getitem__
    array_lookups = []

    def counted_lookup(self: np.lib.npyio.NpzFile, key: str) -> np.ndarray:
        if key != "manifest_json":
            array_lookups.append(key)
        return original(self, key)

    monkeypatch.setattr(np.lib.npyio.NpzFile, "__getitem__", counted_lookup)
    assert EpisodeVideoManifest.validate_file_sidecar(path, _spec(), prepare_cache=prepare_cache)
    assert array_lookups == []


@pytest.mark.parametrize("prepare_cache", [False, True])
def test_sidecar_rejects_object_arrays_without_publication(tmp_path: Path, prepare_cache: bool) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    with np.load(path, allow_pickle=False) as data:
        arrays = {key: data[key] for key in data.files}
    arrays["0/sample_pts"] = np.array([object()], dtype=object)
    np.savez_compressed(path, **arrays)
    assert not EpisodeVideoManifest.validate_file_sidecar(path, _spec(), prepare_cache=prepare_cache)
    assert not list((tmp_path / "cache-home").glob("**/*.bin"))


@pytest.mark.parametrize("published", [False, True])
def test_resolved_sidecar_decompresses_each_array_once(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, published: bool
) -> None:
    original = sidecar_utils._read_arrays
    reads = []

    def counted_read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        reads.append(index)
        return original(archive, index, item)

    def download(path: Path, spec: SidecarSpec) -> bool:
        _write_valid(path, spec)
        return True

    monkeypatch.setattr(sidecar_utils, "_read_arrays", counted_read)
    path = ensure_mp4_sidecar(_spec(), tmp_path, build=_write_valid, download=download if published else None)
    assert EpisodeVideoManifest.validate_file_sidecar(path, _spec())
    records = EpisodeVideoManifest.load_file_sidecar(path)
    np.testing.assert_array_equal(next(iter(records.values())).mp4.sample_pts, [0.0])
    assert reads == [0]
    assert len(list((tmp_path / "cache-home").glob("**/*.bin"))) == 1


def test_invalid_build_preserves_previous_source_and_mapped_generation(tmp_path: Path) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    _write_valid(path, _spec("older"))
    old_bytes = path.read_bytes()
    old = EpisodeVideoManifest.load_file_sidecar(path)

    def invalid_build(target: Path, target_spec: SidecarSpec) -> None:
        _write_valid(target, target_spec)
        with np.load(target, allow_pickle=False) as data:
            arrays = {key: data[key] for key in data.files if key != "0/sample_offsets"}
        np.savez_compressed(target, **arrays)

    with pytest.raises(ValueError, match="failed revision and source validation"):
        ensure_mp4_sidecar(spec, tmp_path, build=invalid_build)
    assert path.read_bytes() == old_bytes
    np.testing.assert_array_equal(next(iter(old.values())).mp4.sample_pts, [0.0])
    assert len(list((tmp_path / "cache-home").glob("**/*.bin"))) == 1
    assert not list((tmp_path / "cache-home").glob("**/*.index.tmp"))


def test_parallel_mapping_protects_legacy_zip_relative_seeks(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Older Python ZIP readers must not seek relative to another worker's position."""
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record(f"{i}.mp4") for i in range(12)], spec=_spec())
    serial, _ = sidecar_utils.mapped_sidecar(path, workers=1)
    original_seek = zipfile._SharedFile.seek
    relative_seeks = []

    def legacy_seek(shared: Any, offset: int, whence: int = os.SEEK_SET) -> int:
        if whence != os.SEEK_CUR:
            return original_seek(shared, offset, whence)
        if threading.current_thread().name.startswith("sidecar-index"):
            assert shared._lock._is_owned(), "ZIP header reads must hold the shared source lock"
            relative_seeks.append(offset)
        # Python 3.12.3 seeks relative to the descriptor, not the member's saved
        # position. Holding the source lock across open prevents interleaving.
        with shared._lock:
            shared._file.seek(offset, whence)
            shared._pos = shared._file.tell()
            return shared._pos

    monkeypatch.setattr(zipfile._SharedFile, "seek", legacy_seek)
    monkeypatch.setenv("HF_LEROBOT_HOME", str(tmp_path / "parallel"))
    parallel, _ = sidecar_utils.mapped_sidecar(path, workers=4)
    assert relative_seeks
    assert parallel.read_bytes() == serial.read_bytes()


def test_parallel_mapping_decompresses_outside_source_lock(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Protect header reads without serializing independent array decompression."""
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record(f"{i}.mp4") for i in range(2)], spec=_spec())
    serial, _ = sidecar_utils.mapped_sidecar(path, workers=1)
    original_read = np.lib.format.read_array
    overlap = threading.Barrier(2)

    def read_array(member: Any, *args: Any, **kwargs: Any) -> np.ndarray:
        if getattr(member, "name", None) in {"0/sample_pts.npy", "1/sample_pts.npy"}:
            overlap.wait(timeout=5)
        return original_read(member, *args, **kwargs)

    monkeypatch.setattr(np.lib.format, "read_array", read_array)
    monkeypatch.setenv("HF_LEROBOT_HOME", str(tmp_path / "parallel"))
    parallel, _ = sidecar_utils.mapped_sidecar(path, workers=2)
    assert parallel.read_bytes() == serial.read_bytes()


def test_parallel_mapping_is_ordered_and_matches_serial_bytes(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = tmp_path / "index.npz"
    records = [_record(f"{index}.mp4") for index in range(12)]
    for index, record in enumerate(records):
        record.mp4.sample_pts[0] = index
    EpisodeVideoManifest.save_file_sidecar(path, records, spec=_spec())
    serial, _ = sidecar_utils.mapped_sidecar(path, workers=1)
    original = sidecar_utils._read_arrays
    second_started = threading.Event()

    def reversed_read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        if index == 0:
            assert second_started.wait(timeout=5), "Decompression did not overlap"
        elif index == 1:
            second_started.set()
        return original(archive, index, item)

    monkeypatch.setattr(sidecar_utils, "_read_arrays", reversed_read)
    monkeypatch.setenv("HF_LEROBOT_HOME", str(tmp_path / "parallel"))
    parallel, _ = sidecar_utils.mapped_sidecar(path, workers=4)
    assert parallel.read_bytes() == serial.read_bytes()
    mapped = EpisodeVideoManifest.load_file_sidecar(path)
    for index in range(12):
        np.testing.assert_array_equal(mapped[f"{index}.mp4"].mp4.sample_pts, [index])


@pytest.mark.parametrize("byte_limit", [1, 2048])
def test_parallel_read_ahead_obeys_count_and_byte_bounds(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, byte_limit: int
) -> None:
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record(f"{i}.mp4") for i in range(12)], spec=_spec())
    payload = sidecar_utils.sidecar_payload(path)
    original = sidecar_utils._read_arrays
    started = []
    lock = threading.Lock()

    def counted_read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        with lock:
            started.append(index)
        return original(archive, index, item)

    monkeypatch.setattr(sidecar_utils, "_read_arrays", counted_read)
    with ZipFile(path) as archive:
        size = sum(archive.getinfo(f"0/{name}.npy").file_size for name in sidecar_utils.ARRAY_NAMES)
        iterator = sidecar_utils._iter_arrays(
            archive, payload["files"], workers=4, max_pending_bytes=byte_limit
        )
        first = next(iterator)
        np.testing.assert_array_equal(first["sample_pts"], [0.0])
        iterator.close()
    assert len(started) <= min(4, max(1, byte_limit // size))


def test_parallel_failure_drains_workers_and_removes_partial_index(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record(f"{i}.mp4") for i in range(12)], spec=_spec())
    original = sidecar_utils._read_arrays
    active = 0
    lock = threading.Lock()

    def failed_read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        nonlocal active
        with lock:
            active += 1
        try:
            if index == 2:
                raise ValueError("injected decompression failure")
            return original(archive, index, item)
        finally:
            with lock:
                active -= 1

    monkeypatch.setattr(sidecar_utils, "_read_arrays", failed_read)
    with pytest.raises(ValueError, match="injected decompression failure"):
        sidecar_utils.mapped_sidecar(path, workers=4)
    assert active == 0
    assert not list((tmp_path / "cache-home").glob("**/*.bin"))
    assert not list((tmp_path / "cache-home").glob("**/*.index.tmp"))


def test_mapping_rejects_invalid_worker_count(tmp_path: Path) -> None:
    path = tmp_path / "index.npz"
    _write_valid(path, _spec())
    with pytest.raises(ValueError, match="workers"):
        sidecar_utils.mapped_sidecar(path, workers=0)


def test_parallel_writer_releases_oversized_record_before_next_read(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    path = tmp_path / "index.npz"
    EpisodeVideoManifest.save_file_sidecar(path, [_record(f"{i}.mp4") for i in range(3)], spec=_spec())
    original_read = sidecar_utils._read_arrays
    original_iter = sidecar_utils._iter_arrays
    references = []

    def read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        assert all(reference() is None for reference in references), "Previous oversized record retained"
        arrays = original_read(archive, index, item)
        references.append(weakref.ref(arrays["sample_pts"]))
        return arrays

    def small_budget(
        archive: ZipFile, files: list[dict[str, Any]], *, workers: int
    ) -> Iterator[dict[str, np.ndarray]]:
        return original_iter(archive, files, workers=workers, max_pending_bytes=1)

    monkeypatch.setattr(sidecar_utils, "_read_arrays", read)
    monkeypatch.setattr(sidecar_utils, "_iter_arrays", small_budget)
    sidecar_utils.mapped_sidecar(path, workers=4)
    assert all(reference() is None for reference in references)


def test_parallel_write_failure_preserves_old_index_and_drains_workers(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    spec = _spec()
    path = sidecar_cache_path(tmp_path, spec)
    _write_valid(path, _spec("old"))
    old_bytes = path.read_bytes()
    old_records = EpisodeVideoManifest.load_file_sidecar(path)
    original = sidecar_utils._read_arrays
    active = 0
    lock = threading.Lock()

    class FailedWriteArray(np.ndarray):
        def tobytes(self, order: str = "C") -> bytes:
            raise OSError(errno.ENOSPC, "injected write failure")

    def read(archive: ZipFile, index: int, item: dict[str, Any]) -> dict[str, np.ndarray]:
        nonlocal active
        with lock:
            active += 1
        try:
            arrays = original(archive, index, item)
            if index == 0:
                arrays["sample_pts"] = arrays["sample_pts"].view(FailedWriteArray)
            return arrays
        finally:
            with lock:
                active -= 1

    def build(target: Path, target_spec: SidecarSpec) -> None:
        records = [_record(f"{index}.mp4") for index in range(12)]
        EpisodeVideoManifest.save_file_sidecar(target, records, spec=target_spec)

    # Use a matching 12-file specification so validation reaches the actual writer.
    spec = replace(spec, source_files=tuple((f"{i}.mp4", 128) for i in range(12)))
    new_path = sidecar_cache_path(tmp_path, spec)
    assert new_path == path
    monkeypatch.setattr(sidecar_utils, "_read_arrays", read)
    with pytest.raises(OSError, match="free space"):
        ensure_mp4_sidecar(spec, tmp_path, build=build)
    assert active == 0
    assert new_path.read_bytes() == old_bytes
    np.testing.assert_array_equal(next(iter(old_records.values())).mp4.sample_pts, [0.0])
    assert len(list((tmp_path / "cache-home").glob("**/*.bin"))) == 1
    assert not list((tmp_path / "cache-home").glob("**/*.index.tmp"))
    assert not list(new_path.parent.glob("*.tmp.npz"))
