"""Token shard reader. Shards are flat uint16 .bin files written by scripts/prepare_data.py."""

import hashlib
import json
import os
import struct

import numpy as np


RUN1_LEGACY_NUMPY_VERSION = "2.5.1"
NO_FIM_DERIVATION_ALGORITHM = (
    "run1_deterministic_no_fim_normalization_v1"
)
DERIVED_KIND_CODES = {"l2r": 0, "fim_psm": 1, "fim_spm": 2}
DERIVED_UNIT_RECORD_DTYPE = np.dtype(
    [
        ("length", "<u4"),
        ("kind", "u1"),
        ("reserved", "u1", (3,)),
    ]
)


def canonical_json_sha256(value) -> str:
    rendered = json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )
    return hashlib.sha256(rendered.encode("utf-8")).hexdigest()


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


def stable_file_sha256(path: str) -> str:
    """Hash one regular file and reject replacement or mutation during read."""
    if os.path.islink(path):
        raise ValueError(f"file cannot be a symbolic link: {path}")
    before = os.stat(path, follow_symlinks=False)
    if not os.path.isfile(path):
        raise ValueError(f"path is not a regular file: {path}")
    digest = file_sha256(path)
    after = os.stat(path, follow_symlinks=False)
    identity = lambda value: (
        value.st_dev,
        value.st_ino,
        value.st_size,
        value.st_mtime_ns,
    )
    if identity(before) != identity(after):
        raise ValueError(f"file changed during SHA-256 read: {path}")
    return digest


def valid_sha256(value) -> bool:
    return (
        isinstance(value, str)
        and len(value) == 64
        and all(char in "0123456789abcdef" for char in value)
    )


def load_json_object(path):
    with open(path, encoding="utf-8") as f:
        value = json.load(f)
    if not isinstance(value, dict):
        raise ValueError(f"{path}: top-level JSON must be an object")
    return value


def validate_data_integrity_reference(
    config,
    index,
    index_path,
    problems,
):
    """Validate the run 1 sidecar and return hashes for current shards."""
    contract = config.get("data_integrity")
    if contract is None:
        return {}, None
    if not isinstance(contract, dict):
        problems.append("config data_integrity is not an object")
        return {}, None
    role = contract.get("role")
    if role not in ("current", "source"):
        problems.append("config data_integrity role is not current or source")
    receipt_path = contract.get("receipt")
    receipt_sha256 = contract.get("sha256")
    if not isinstance(receipt_path, str) or not os.path.isfile(receipt_path):
        problems.append("config data_integrity receipt is missing")
        return {}, None
    if not valid_sha256(receipt_sha256):
        problems.append("config data_integrity receipt SHA-256 is malformed")
        return {}, None
    try:
        actual_receipt_sha256 = stable_file_sha256(receipt_path)
    except (OSError, ValueError) as exc:
        problems.append(f"data integrity receipt cannot be hashed: {exc}")
        return {}, None
    if actual_receipt_sha256 != receipt_sha256:
        problems.append("data integrity receipt hash differs from config")
        return {}, None
    try:
        receipt = load_json_object(receipt_path)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        problems.append(f"data integrity receipt cannot be loaded: {exc}")
        return {}, None
    if (
        receipt.get("schema_version") != 1
        or receipt.get("status") != "complete"
        or receipt.get("algorithm") != NO_FIM_DERIVATION_ALGORITHM
    ):
        problems.append("data integrity receipt contract differs")
    tokenizer = receipt.get("tokenizer", {})
    tokenizer_path = config.get("tokenizer_path")
    if (
        not isinstance(tokenizer_path, str)
        or not os.path.isfile(tokenizer_path)
        or tokenizer.get("sha256") != file_sha256(tokenizer_path)
    ):
        problems.append("data integrity tokenizer differs from config")

    if role == "source":
        return {}, receipt

    source_index = receipt.get("source_index", {})
    if (
        os.path.abspath(source_index.get("path", ""))
        != os.path.abspath(index_path)
        or source_index.get("sha256") != file_sha256(index_path)
    ):
        problems.append("current data integrity source index differs")
    expected_hashes = {}
    for split in ("train", "val"):
        entries = index.get("splits", {}).get(split, [])
        split_receipt = receipt.get("splits", {}).get(split, {})
        receipt_entries = split_receipt.get("shards", [])
        core = [
            {
                key: entry.get(key)
                for key in ("path", "tokens", "bytes", "sha256")
            }
            for entry in receipt_entries
            if isinstance(entry, dict)
        ]
        if canonical_json_sha256(core) != split_receipt.get(
            "shard_manifest_sha256"
        ):
            problems.append(
                f"data integrity split {split} manifest hash differs"
            )
        if [entry.get("path") for entry in entries] != [
            entry.get("path") for entry in receipt_entries
        ]:
            problems.append(
                f"data integrity split {split} shard order differs"
            )
            continue
        if split_receipt.get("source_tokens") != sum(
            entry.get("tokens", 0)
            for entry in entries
            if isinstance(entry, dict)
        ):
            problems.append(
                f"data integrity split {split} token total differs"
            )
        for index_entry, receipt_entry in zip(entries, receipt_entries):
            if (
                receipt_entry.get("tokens") != index_entry.get("tokens")
                or receipt_entry.get("bytes")
                != index_entry.get("tokens", 0) * 2
                or not valid_sha256(receipt_entry.get("sha256"))
            ):
                problems.append(
                    f"data integrity split {split} shard evidence differs"
                )
                continue
            expected_hashes[index_entry["path"]] = receipt_entry["sha256"]
    return expected_hashes, receipt


class _DerivedEvidence:
    """Independent destination-unit hasher for schema-3 data indexes."""

    def __init__(self, eos_token_id, domain):
        self.eos_token_id = eos_token_id
        self.domain = domain.encode("utf-8")
        self.ordered = hashlib.sha256()
        self.destination = hashlib.sha256()
        for digest, label in (
            (self.ordered, b"RECOVERED_RAW"),
            (self.destination, b"DESTINATION"),
        ):
            digest.update(b"WISP_NO_FIM_V1\0" + label + b"\0")
            digest.update(struct.pack("<Q", len(self.domain)))
            digest.update(self.domain)
        self.units = 0
        self.source_tokens = 0
        self.raw_tokens = 0
        self.derived_tokens = 0
        self.counts = {"l2r": 0, "fim_psm": 0, "fim_spm": 0}
        self.eos_bytes = np.asarray([eos_token_id], dtype="<u2").tobytes()

    def add(self, kind, raw):
        if kind not in DERIVED_KIND_CODES:
            raise ValueError(f"unknown derived unit kind: {kind}")
        raw = np.asarray(raw, dtype="<u2")
        if kind == "l2r":
            if not 1 <= raw.size <= 1024:
                raise ValueError("derived plain unit length differs")
        elif not 16 <= raw.size <= 1024:
            raise ValueError("derived FIM unit length differs")
        kind_code = DERIVED_KIND_CODES[kind]
        raw_bytes = raw.tobytes()
        ordinal = self.units
        self.ordered.update(
            b"U"
            + struct.pack("<QBQ", ordinal, kind_code, int(raw.size))
        )
        self.ordered.update(raw_bytes)
        self.destination.update(
            b"U"
            + struct.pack(
                "<QBQ",
                ordinal,
                kind_code,
                int(raw.size) + 1,
            )
        )
        self.destination.update(raw_bytes)
        self.destination.update(self.eos_bytes)
        self.units += 1
        self.raw_tokens += int(raw.size)
        self.derived_tokens += int(raw.size) + 1
        self.source_tokens += int(raw.size) + (
            4 if kind != "l2r" else 1
        )
        self.counts[kind] += 1

    def summary(self):
        fim_units = self.counts["fim_psm"] + self.counts["fim_spm"]
        footer = b"END" + struct.pack(
            "<QQQQQQ",
            self.units,
            self.source_tokens,
            self.raw_tokens,
            self.derived_tokens,
            self.counts["l2r"],
            fim_units,
        )
        ordered = self.ordered.copy()
        ordered.update(footer)
        destination = self.destination.copy()
        destination.update(footer)
        return {
            "units": self.units,
            "source_tokens": self.source_tokens,
            "raw_tokens": self.raw_tokens,
            "derived_tokens": self.derived_tokens,
            "removed_fim_tokens": self.source_tokens - self.derived_tokens,
            "l2r_units": self.counts["l2r"],
            "fim_psm_units": self.counts["fim_psm"],
            "fim_spm_units": self.counts["fim_spm"],
            "ordered_raw_units_sha256": ordered.hexdigest(),
            "destination_units_sha256": destination.hexdigest(),
        }


def _normalization_evidence(value):
    keys = (
        "units",
        "source_tokens",
        "raw_tokens",
        "derived_tokens",
        "removed_fim_tokens",
        "l2r_units",
        "fim_psm_units",
        "fim_spm_units",
        "ordered_raw_units_sha256",
        "destination_units_sha256",
    )
    return {key: value.get(key) for key in keys}


def validate_derived_unit_sidecars(
    index,
    actual_dir,
    tokenizer_sha256,
    problems,
):
    """Independently verify schema-3 unit boundaries and normalization hashes."""
    build = index.get("build", {})
    special = build.get("special_token_ids", {})
    if (
        set(special) != {"eos", "prefix", "middle", "suffix"}
        or any(
            not isinstance(value, int) or isinstance(value, bool)
            for value in special.values()
        )
        or len(set(special.values())) != 4
    ):
        problems.append("derived special-token ids are malformed")
        return
    fim_ids = [special[key] for key in ("prefix", "middle", "suffix")]
    inverse_kinds = {
        value: key for key, value in DERIVED_KIND_CODES.items()
    }
    seen_sidecars = set()
    for split in ("train", "val"):
        split_evidence = _DerivedEvidence(
            special["eos"],
            f"split:{split}:tokenizer:{tokenizer_sha256}",
        )
        entries = index.get("splits", {}).get(split, [])
        for entry in entries:
            sidecar = entry.get("unit_sidecar", {})
            relative = sidecar.get("path")
            valid_path = (
                isinstance(relative, str)
                and relative
                and not os.path.isabs(relative)
                and os.path.normpath(relative) == relative
                and relative != ".."
                and not relative.startswith(".." + os.sep)
            )
            if not valid_path or relative in seen_sidecars:
                problems.append(
                    f"derived split {split} unit sidecar path is unsafe"
                )
                continue
            seen_sidecars.add(relative)
            path = os.path.join(actual_dir, entry["path"])
            sidecar_path = os.path.join(actual_dir, relative)
            if os.path.islink(sidecar_path) or not os.path.isfile(
                sidecar_path
            ):
                problems.append(
                    f"derived unit sidecar is missing: {relative}"
                )
                continue
            if (
                sidecar.get("record_format")
                != "uint32_length_uint8_kind_3_zero_bytes"
                or sidecar.get("bytes")
                != sidecar.get("records", 0)
                * DERIVED_UNIT_RECORD_DTYPE.itemsize
                or os.path.getsize(sidecar_path) != sidecar.get("bytes")
                or not valid_sha256(sidecar.get("sha256"))
            ):
                problems.append(
                    f"derived unit sidecar declaration differs: {relative}"
                )
                continue
            try:
                if stable_file_sha256(sidecar_path) != sidecar["sha256"]:
                    problems.append(
                        f"derived unit sidecar SHA-256 differs: {relative}"
                    )
                    continue
                values = np.memmap(path, dtype="<u2", mode="r")
                records = np.memmap(
                    sidecar_path,
                    dtype=DERIVED_UNIT_RECORD_DTYPE,
                    mode="r",
                )
                if np.any(records["reserved"] != 0):
                    raise ValueError("reserved sidecar bytes are not zero")
                if int(records["length"].sum(dtype=np.uint64)) != int(
                    values.size
                ):
                    raise ValueError("unit lengths do not cover shard")
                if np.any(records["kind"] > max(inverse_kinds)):
                    raise ValueError("unit kind is outside registered values")
                for token_id in fim_ids:
                    if np.any(values == token_id):
                        raise ValueError("destination retains a FIM sentinel")
                shard_evidence = _DerivedEvidence(
                    special["eos"],
                    (
                        f"shard:{split}:{entry['path']}:"
                        f"tokenizer:{tokenizer_sha256}"
                    ),
                )
                offset = 0
                for record in records:
                    length = int(record["length"])
                    end = offset + length
                    if length < 2 or end > values.size:
                        raise ValueError("unit boundary is outside shard")
                    unit = np.asarray(values[offset:end])
                    if int(unit[-1]) != special["eos"]:
                        raise ValueError("derived unit lacks writer EOS")
                    kind = inverse_kinds[int(record["kind"])]
                    raw = unit[:-1]
                    shard_evidence.add(kind, raw)
                    split_evidence.add(kind, raw)
                    offset = end
                if offset != values.size:
                    raise ValueError("unit boundaries do not cover shard")
            except (OSError, ValueError) as exc:
                problems.append(
                    f"derived unit sidecar verification failed for "
                    f"{relative}: {exc}"
                )
                continue
            actual = shard_evidence.summary()
            if _normalization_evidence(actual) != entry.get(
                "normalization"
            ):
                problems.append(
                    f"derived shard normalization differs: {entry['path']}"
                )
            source = entry.get("source", {})
            if (
                source.get("path") != entry.get("path")
                or source.get("tokens") != actual["source_tokens"]
                or not valid_sha256(source.get("sha256"))
                or not valid_sha256(
                    source.get("source_wire_units_sha256")
                )
                or source.get("ordered_raw_units_sha256")
                != actual["ordered_raw_units_sha256"]
            ):
                problems.append(
                    f"derived source linkage differs: {entry['path']}"
                )
        expected = build.get("split_evidence", {}).get(split, {})
        if _normalization_evidence(split_evidence.summary()) != (
            _normalization_evidence(expected)
        ):
            problems.append(
                f"derived split {split} normalization hash differs"
            )


def validate_data_contract(config: dict, index_path: str) -> dict:
    """Require trainer-visible preprocessing claims to match frozen shards."""
    with open(index_path, encoding="utf-8") as f:
        index = json.load(f)
    problems = []
    actual_index = os.path.abspath(index_path)
    configured_index = config.get("data_index")
    expected_index = (
        os.path.abspath(configured_index)
        if isinstance(configured_index, str)
        else None
    )
    if expected_index != actual_index:
        problems.append(
            f"config data_index {expected_index} != supplied {actual_index}"
        )
    actual_dir = os.path.dirname(actual_index)
    configured_dir = config.get("data_dir")
    expected_dir = (
        os.path.abspath(configured_dir)
        if isinstance(configured_dir, str)
        else None
    )
    if expected_dir != actual_dir:
        problems.append(
            f"config data_dir {expected_dir} != index directory {actual_dir}"
        )
    for key in ("vocab_size", "fim_rate", "fim_chunk"):
        if config.get(key) != index.get(key):
            problems.append(
                f"config {key} {config.get(key)!r} != "
                f"index {key} {index.get(key)!r}"
            )
    if problems:
        expected_integrity_hashes, integrity_receipt = {}, None
    else:
        expected_integrity_hashes, integrity_receipt = (
            validate_data_integrity_reference(
                config,
                index,
                index_path,
                problems,
            )
        )
    build_contract = config.get("data_build_contract")
    build_contract_schema = None
    tokenizer_digest = None
    if build_contract is not None:
        if not isinstance(build_contract, dict):
            problems.append("config data_build_contract is not an object")
            build_contract = {}
        build_contract_schema = build_contract.get("schema_version")
        if build_contract.get("require_fresh_output_dir") is not True:
            problems.append(
                "config data_build_contract does not require a fresh directory"
            )
        build = index.get("build")
        if not isinstance(build, dict):
            problems.append("attested data index build record is missing")
            build = {}
        tokenizer_path = config.get("tokenizer_path")
        if not isinstance(tokenizer_path, str) or not os.path.isfile(
            tokenizer_path
        ):
            problems.append("config tokenizer_path is missing or unreadable")
            tokenizer_digest = None
        else:
            tokenizer_digest = file_sha256(tokenizer_path)
            if build.get("tokenizer_sha256") != tokenizer_digest:
                problems.append(
                    "data build tokenizer hash differs from current tokenizer"
                )
        if build.get("completed") is not True:
            problems.append("data build is not marked complete")

        if build_contract_schema == 1:
            for key in ("train_tokens", "validation_tokens"):
                value = build_contract.get(key)
                if (
                    not isinstance(value, int)
                    or isinstance(value, bool)
                    or value < 1
                ):
                    problems.append(
                        f"config data_build_contract {key} is not positive"
                    )
            if config.get("seed") != build_contract.get("seed"):
                problems.append(
                    f"config seed {config.get('seed')!r} != data build "
                    f"contract seed {build_contract.get('seed')!r}"
                )
            if index.get("schema_version") != 2:
                problems.append("attested data index schema_version is not 2")
            expected_build = {
                "train_tokens_requested": build_contract.get("train_tokens"),
                "validation_tokens_requested": build_contract.get(
                    "validation_tokens"
                ),
                "seed": build_contract.get("seed"),
                "fresh_output_directory": build_contract.get(
                    "require_fresh_output_dir"
                ),
                "config_canonical_sha256": canonical_json_sha256(config),
                "sources_canonical_sha256": canonical_json_sha256(
                    config.get("sources")
                ),
            }
            for key, expected in expected_build.items():
                if build.get(key) != expected:
                    problems.append(
                        f"data build {key} {build.get(key)!r} != "
                        f"config contract {expected!r}"
                    )
        elif build_contract_schema == 2:
            if build_contract.get(
                "strategy"
            ) != NO_FIM_DERIVATION_ALGORITHM:
                problems.append("data derivation strategy differs")
            if "sources" in config:
                problems.append(
                    "derived no-FIM config must not contain executable sources"
                )
            integrity = config.get("data_integrity", {})
            if (
                not isinstance(integrity, dict)
                or integrity.get("role") != "source"
            ):
                problems.append(
                    "derived no-FIM config integrity role is not source"
                )
            expected_contract = {
                "schema_version": 2,
                "strategy": NO_FIM_DERIVATION_ALGORITHM,
                "source_index": (
                    integrity_receipt.get("source_index")
                    if isinstance(integrity_receipt, dict)
                    else None
                ),
                "source_integrity_receipt": (
                    integrity.get("receipt")
                    if isinstance(integrity, dict)
                    else None
                ),
                "source_integrity_receipt_sha256": (
                    integrity.get("sha256")
                    if isinstance(integrity, dict)
                    else None
                ),
                "require_fresh_output_dir": True,
            }
            if build_contract != expected_contract:
                problems.append(
                    "data derivation contract differs from source integrity"
                )
            if index.get("schema_version") != 3:
                problems.append("derived data index schema_version is not 3")
            expected_build = {
                "strategy": NO_FIM_DERIVATION_ALGORITHM,
                "fresh_output_directory": True,
                "config_canonical_sha256": canonical_json_sha256(config),
                "source_index": build_contract.get("source_index"),
                "source_integrity_receipt": {
                    "path": build_contract.get(
                        "source_integrity_receipt"
                    ),
                    "sha256": build_contract.get(
                        "source_integrity_receipt_sha256"
                    ),
                },
                "contract": build_contract,
            }
            for key, expected in expected_build.items():
                if build.get(key) != expected:
                    problems.append(
                        f"derived data build {key} differs from config"
                    )
            script = build.get("derivation_script", {})
            if (
                not isinstance(script, dict)
                or not isinstance(script.get("path"), str)
                or not os.path.isfile(script["path"])
                or not valid_sha256(script.get("sha256"))
                or file_sha256(script["path"]) != script["sha256"]
            ):
                problems.append("data derivation script hash differs")
        else:
            problems.append(
                "config data_build_contract schema is not 1 or 2"
            )
    splits = index.get("splits")
    if not isinstance(splits, dict):
        problems.append("index splits object is missing")
        splits = {}
    seen_paths = set()
    for split in ("train", "val"):
        entries = splits.get(split)
        if not isinstance(entries, list) or not entries:
            problems.append(f"index split {split} is empty or missing")
            continue
        if any(not isinstance(entry, dict) for entry in entries):
            problems.append(f"index split {split} has a non-object shard entry")
            continue
        declared_totals = [entry.get("total_tokens") for entry in entries]
        token_counts = [entry.get("tokens") for entry in entries]
        valid_declared_totals = all(
            isinstance(value, int)
            and not isinstance(value, bool)
            and value > 0
            for value in declared_totals
        )
        valid_token_counts = all(
            isinstance(value, int)
            and not isinstance(value, bool)
            and value >= 1
            for value in token_counts
        )
        if (
            not valid_declared_totals
            or not valid_token_counts
            or len(set(declared_totals)) != 1
            or sum(token_counts) != declared_totals[0]
        ):
            problems.append(f"index split {split} token totals are inconsistent")
        for entry in entries:
            relative = entry.get("path")
            valid_path = (
                isinstance(relative, str)
                and relative
                and not os.path.isabs(relative)
                and os.path.normpath(relative) == relative
                and relative != ".."
                and not relative.startswith(".." + os.sep)
            )
            if not valid_path:
                problems.append(
                    f"index split {split} has unsafe shard path {relative!r}"
                )
                continue
            if relative in seen_paths:
                problems.append(
                    f"index repeats shard path across splits: {relative}"
                )
                continue
            seen_paths.add(relative)
            shard_path = os.path.join(actual_dir, relative)
            tokens = entry.get("tokens")
            if os.path.islink(shard_path):
                problems.append(
                    f"index split {split} shard is a symbolic link: {relative}"
                )
            elif not os.path.isfile(shard_path):
                problems.append(
                    f"index split {split} shard is missing: {relative}"
                )
            elif (
                isinstance(tokens, int)
                and not isinstance(tokens, bool)
                and os.path.getsize(shard_path) != tokens * 2
            ):
                problems.append(
                    f"index split {split} shard byte size differs: {relative}"
                )
            else:
                expected_hash = expected_integrity_hashes.get(relative)
                hash_bound_index = (
                    index.get("schema_version") == 3
                    or (
                        index.get("schema_version") == 2
                        and build_contract_schema == 1
                    )
                )
                if hash_bound_index:
                    declared_bytes = entry.get("bytes")
                    declared_hash = entry.get("sha256")
                    expected_bytes = (
                        tokens * 2
                        if isinstance(tokens, int)
                        and not isinstance(tokens, bool)
                        else None
                    )
                    if (
                        expected_bytes is None
                        or declared_bytes != expected_bytes
                    ):
                        problems.append(
                            f"attested shard byte declaration differs: {relative}"
                        )
                    if not valid_sha256(declared_hash):
                        problems.append(
                            f"attested shard SHA-256 is malformed: {relative}"
                        )
                    elif (
                        expected_hash is not None
                        and expected_hash != declared_hash
                    ):
                        problems.append(
                            f"attested and sidecar shard hashes differ: {relative}"
                        )
                    expected_hash = declared_hash
                if expected_hash is not None and valid_sha256(expected_hash):
                    try:
                        actual_hash = stable_file_sha256(shard_path)
                    except (OSError, ValueError) as exc:
                        problems.append(
                            f"shard integrity read failed for {relative}: {exc}"
                        )
                    else:
                        if actual_hash != expected_hash:
                            problems.append(
                                f"shard SHA-256 differs: {relative}"
                            )
        if index.get("schema_version") == 3:
            evidence = index.get("build", {}).get(
                "split_evidence", {}
            ).get(split, {})
            total = (
                declared_totals[0]
                if valid_declared_totals and len(set(declared_totals)) == 1
                else None
            )
            if evidence.get("derived_tokens") != total:
                problems.append(
                    f"derived split {split} evidence token total differs"
                )
            fim_units = evidence.get("fim_psm_units", 0) + evidence.get(
                "fim_spm_units", 0
            )
            if (
                not isinstance(fim_units, int)
                or evidence.get("removed_fim_tokens") != 3 * fim_units
                or evidence.get("source_tokens", 0)
                - evidence.get("derived_tokens", 0)
                != evidence.get("removed_fim_tokens")
            ):
                problems.append(
                    f"derived split {split} FIM removal arithmetic differs"
                )
            if evidence.get("units") != (
                evidence.get("l2r_units", 0) + fim_units
            ):
                problems.append(
                    f"derived split {split} unit counts differ"
                )
            for key in (
                "ordered_raw_units_sha256",
                "destination_units_sha256",
            ):
                if not valid_sha256(evidence.get(key)):
                    problems.append(
                        f"derived split {split} {key} is malformed"
                    )
    if index.get("schema_version") == 3 and isinstance(
        tokenizer_digest, str
    ):
        validate_derived_unit_sidecars(
            index,
            actual_dir,
            tokenizer_digest,
            problems,
        )
    if (
        build_contract_schema == 1
        and isinstance(splits, dict)
    ):
        expected_totals = {
            "train": build_contract.get("train_tokens"),
            "val": build_contract.get("validation_tokens"),
        }
        for split, expected in expected_totals.items():
            entries = splits.get(split)
            if (
                isinstance(entries, list)
                and entries
                and isinstance(entries[0], dict)
                and isinstance(entries[0].get("total_tokens"), int)
                and isinstance(expected, int)
                and entries[0]["total_tokens"] < expected
            ):
                problems.append(
                    f"index split {split} has fewer tokens than requested"
                )
    if problems:
        raise ValueError(
            "training data contract mismatch:\n- "
            + "\n- ".join(problems)
        )
    return index


def sampler_reset_steps(config: dict) -> list[int]:
    """Validate registered sampler resets caused by known process recovery."""
    value = config.get("sampler_reset_steps", [])
    valid = (
        isinstance(value, list)
        and all(
            isinstance(step, int)
            and not isinstance(step, bool)
            and 0 < step < config["max_steps"]
            for step in value
        )
        and value == sorted(set(value))
    )
    if not valid:
        raise ValueError(
            "sampler_reset_steps must be sorted unique integers greater "
            "than zero and less than max_steps"
        )
    return value


def sampler_batches_since_reset(
    completed_steps: int,
    grad_accum: int,
    reset_steps: list[int],
) -> int:
    """Count RNG batches after the latest reset applied before this boundary."""
    prior_resets = [step for step in reset_steps if step < completed_steps]
    latest_reset = max(prior_resets, default=0)
    return (completed_steps - latest_reset) * grad_accum


def validate_resume_sampling_contract(checkpoint_meta: dict, config: dict):
    """Reject sampling changes across resume, with one recorded run1 exception."""
    checkpoint_config = checkpoint_meta.get("config")
    if not isinstance(checkpoint_config, dict):
        raise ValueError("checkpoint config is missing")
    fields = (
        "run_name",
        "data_index",
        "seed",
        "seq_len",
        "mtp_depth",
        "micro_batch",
        "grad_accum",
    )
    for field in fields:
        if checkpoint_config.get(field) != config.get(field):
            raise ValueError(
                f"resume sampling field {field} "
                f"{config.get(field)!r} != checkpoint "
                f"{checkpoint_config.get(field)!r}"
            )
    checkpoint_resets = checkpoint_config.get("sampler_reset_steps", [])
    current_resets = config.get("sampler_reset_steps", [])
    if checkpoint_resets == current_resets:
        if not isinstance(checkpoint_meta.get("train_sampler"), dict):
            raise ValueError(
                "checkpoint is missing exact training sampler state"
            )
        return {"legacy_reset_registration": False}
    legacy_reset_registration = (
        checkpoint_config.get("run_name") == "wisp-run1-110m-code"
        and config.get("run_name") == "wisp-run1-110m-code"
        and "sampler_reset_steps" not in checkpoint_config
        and current_resets == [300]
        and isinstance(checkpoint_meta.get("step"), int)
        and not isinstance(checkpoint_meta.get("step"), bool)
        and checkpoint_meta["step"] >= 300
        and checkpoint_meta.get("train_sampler") is None
        and np.__version__ == RUN1_LEGACY_NUMPY_VERSION
    )
    if not legacy_reset_registration:
        raise ValueError(
            "resume sampler reset schedule differs from checkpoint lineage"
        )
    return {"legacy_reset_registration": True}


class ShardDataset:
    """
    Random-offset sampler over a set of memory-mapped uint16 token shards.

    Documents are already concatenated with an EOS separator at prepare time, so a
    random window is a valid training example. Windows are `span` tokens long,
    where span = seq_len + 1 + mtp_depth.
    """

    def __init__(self, index_path: str, split: str, span: int, seed: int = 1337):
        with open(index_path) as f:
            index = json.load(f)
        if split not in index["splits"]:
            raise KeyError(f"split {split!r} not in {list(index['splits'])}")

        root = os.path.dirname(os.path.abspath(index_path))
        self.index_path = os.path.abspath(index_path)
        self.index_sha256 = file_sha256(self.index_path)
        self.split = split
        self.span = span
        self.seed = seed
        self.shards = []
        self.lengths = []
        for entry in index["splits"][split]:
            path = os.path.join(root, entry["path"])
            arr = np.memmap(path, dtype=np.uint16, mode="r")
            if arr.shape[0] <= span:
                continue
            self.shards.append(arr)
            self.lengths.append(arr.shape[0] - span)
        if not self.shards:
            raise RuntimeError(f"no usable shards for split {split!r}")

        self.total = int(sum(self.lengths))
        self.weights = np.array(self.lengths, dtype=np.float64) / self.total
        self.rng = np.random.default_rng(seed)
        self.batches_drawn = 0
        self.vocab_size = index["vocab_size"]
        self.token_count = int(index["splits"][split][0].get("total_tokens", 0)) or None

    def __len__(self) -> int:
        return self.total

    def _draw_coordinates_from_rng(
        self,
        rng: np.random.Generator,
        batch_size: int,
    ):
        shard_ids = rng.choice(
            len(self.shards),
            size=batch_size,
            p=self.weights,
        )
        starts = np.empty(batch_size, dtype=np.int64)
        for row, sid in enumerate(shard_ids):
            starts[row] = rng.integers(0, self.lengths[sid])
        return shard_ids, starts

    def _draw_coordinates(self, batch_size: int):
        shard_ids, starts = self._draw_coordinates_from_rng(
            self.rng,
            batch_size,
        )
        self.batches_drawn += 1
        return shard_ids, starts

    def batch(self, batch_size: int) -> np.ndarray:
        """Returns an (batch_size, span) int32 array."""
        out = np.empty((batch_size, self.span), dtype=np.int32)
        shard_ids, starts = self._draw_coordinates(batch_size)
        for row, (sid, start) in enumerate(zip(shard_ids, starts)):
            out[row] = self.shards[sid][start:start + self.span].astype(np.int32)
        return out

    def reset_sampler(self):
        """Reset to the registered seed, matching a fresh process exactly."""
        self.rng = np.random.default_rng(self.seed)
        self.batches_drawn = 0

    def advance_batches(self, batch_size: int, batches: int):
        """Reconstruct a legacy checkpoint's RNG state without reading tokens."""
        if (
            not isinstance(batches, int)
            or isinstance(batches, bool)
            or batches < 0
        ):
            raise ValueError("batches to advance must be a non-negative integer")
        if self.batches_drawn != 0:
            raise ValueError("sampler can only advance from its initial state")
        for _ in range(batches):
            self._draw_coordinates(batch_size)

    def sampler_state(self, batch_size: int) -> dict:
        """Return a JSON-serializable exact training-sampler checkpoint."""
        rng_state = self.rng.bit_generator.state
        return {
            "schema_version": 1,
            "index_path": self.index_path,
            "index_sha256": self.index_sha256,
            "split": self.split,
            "span": self.span,
            "seed": self.seed,
            "batch_size": batch_size,
            "batches_drawn_since_reset": self.batches_drawn,
            "bit_generator": type(self.rng.bit_generator).__name__,
            "numpy_version": np.__version__,
            "rng_state": rng_state,
            "rng_state_sha256": canonical_json_sha256(rng_state),
        }

    def restore_sampler_state(
        self,
        state: dict,
        batch_size: int,
        expected_batches: int,
    ):
        """Restore and validate an exact training-sampler checkpoint."""
        if not isinstance(state, dict) or state.get("schema_version") != 1:
            raise ValueError("training sampler state schema is not 1")
        expected = {
            "index_sha256": self.index_sha256,
            "split": self.split,
            "span": self.span,
            "seed": self.seed,
            "batch_size": batch_size,
            "batches_drawn_since_reset": expected_batches,
            "bit_generator": type(self.rng.bit_generator).__name__,
            "numpy_version": np.__version__,
        }
        for key, value in expected.items():
            if state.get(key) != value:
                raise ValueError(
                    f"training sampler state {key} {state.get(key)!r} "
                    f"!= expected {value!r}"
                )
        rng_state = state.get("rng_state")
        if not isinstance(rng_state, dict):
            raise ValueError("training sampler RNG state is missing")
        if state.get("rng_state_sha256") != canonical_json_sha256(rng_state):
            raise ValueError("training sampler RNG state hash differs")
        candidate_rng = np.random.default_rng(self.seed)
        try:
            candidate_rng.bit_generator.state = rng_state
        except (TypeError, ValueError) as exc:
            raise ValueError("training sampler RNG state is invalid") from exc
        expected_rng = np.random.default_rng(self.seed)
        for _ in range(expected_batches):
            self._draw_coordinates_from_rng(expected_rng, batch_size)
        if rng_state != expected_rng.bit_generator.state:
            raise ValueError(
                "training sampler RNG state does not match deterministic replay"
            )
        self.rng.bit_generator.state = rng_state
        self.batches_drawn = expected_batches

    def iter_eval(self, batch_size: int, n_batches: int, seed: int = 7):
        """Deterministic batches for held-out evaluation."""
        rng = np.random.default_rng(seed)
        for _ in range(n_batches):
            out = np.empty((batch_size, self.span), dtype=np.int32)
            shard_ids = rng.choice(len(self.shards), size=batch_size, p=self.weights)
            for row, sid in enumerate(shard_ids):
                start = rng.integers(0, self.lengths[sid])
                out[row] = self.shards[sid][start:start + self.span].astype(np.int32)
            yield out