# Map-style reader over the built mix, designed so that resumption is an index advance, not a replay.

# Why this file exists in this shape (docs/01-plan.md §5, "the resume trap"): `transformers.Trainer` can
# resume an `IterableDataset` only by *decoding and discarding* the batches before the checkpoint step,
# which costs O(steps) of I/O after every interruption -- up to 3,815 steps of nothing. A map-style
# reader over a concatenated token store makes the same thing a one-line offset. §3.1's "same position"
# is therefore a number in the checkpoint, not a feeling.
#
# Sampling policy, stated so it can be reproduced:
#   stream   = the uint16 id arrays of manifest["shards"] in manifest order, concatenated
#   sample i = stream[i*L : (i+1)*L]  with L = seq_len, no padding, documents packed back to back
#   order    = a single global permutation of sample indices from numpy's PCG64(seed)
#
# Packing without a document boundary means a few windows per shard splice two unrelated documents and
# let attention cross the seam. That is what TinyLlama, SmolLM and GPT-2-era pretraining all do; a
# block-diagonal mask would cost activation memory we do not have at bs 2x2048 on a T4. Recorded as a
# known, deliberate deviation rather than a silence.
#
# The permutation is over ~490k indices (1B tokens / 2048), i.e. ~2 MB, so a *global* shuffle is free
# here -- no need for the windowed approximations that streaming readers settle for.
#
# torch is imported lazily-ish: the byte-level part of this module is usable in a CPU-only check.

import bisect
import json
import os
import struct

import numpy as np


def read_shard_index(path):
    """(vocab, n_docs, n_tokens, offset_bytes, ids_bytes) for one shard without loading its ids."""
    with open(path, "rb") as f:
        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", f.read(16))
    if flags != 0:
        raise ValueError(f"{path}: unknown shard flags {flags}")
    return vocab, n_docs, n_tokens, 16 + 4 * (n_docs + 1), 2 * n_tokens


class PackedTokenStore:
    """Random-access uint16 view of the concatenated mix, in manifest order."""

    def __init__(self, root, manifest=None, files=None, vocab=None):
        man = manifest if manifest is not None else json.load(
            open(os.path.join(root, "manifest.json")))
        self.vocab = vocab or man["tokenizer"]["vocab_size"]
        names = files if files is not None else [s["file"] for s in man["shards"]]
        self.shards, self.starts, self._maps = [], [], []
        tot = 0
        for name in names:
            p = name if os.path.isabs(name) else os.path.join(root, name)
            v, n_docs, n_tokens, ob, ib = read_shard_index(p)
            if v != self.vocab:
                raise ValueError(f"{p}: shard vocab {v} != manifest vocab {self.vocab}")
            # memmap, never load: the whole mix is ~2.6 GB and a GPU session has 30 GiB shared with
            # two CUDA contexts. Page faults are the point.
            mm = np.memmap(p, dtype=np.uint16, mode="r", offset=ob, shape=(n_tokens,))
            self.shards.append((p, mm, n_tokens))
            self.starts.append(tot)
            tot += n_tokens
        self.total_tokens = tot
        if tot == 0:
            # Zero-length store must not look like a valid empty dataset (E-017's shape).
            raise ValueError("no tokens in the mix -- nothing was built or nothing was restored")

    def tokens(self, g0, n):
        """Global token span [g0, g0+n) as a contiguous uint16 array, across shard boundaries."""
        out = np.empty(n, dtype=np.uint16)
        done = 0
        while done < n:
            si = bisect.bisect_right(self.starts, g0 + done) - 1
            p, mm, nt = self.shards[si]
            local = (g0 + done) - self.starts[si]
            take = min(n - done, nt - local)
            if take <= 0:
                raise IndexError(f"span [{g0},{g0+n}) runs past the end of the mix ({self.total_tokens})")
            out[done:done + take] = mm[local:local + take]
            done += take
        return out


def perm_sha(n_samples, seed, base=None):
    """Identity of the visit order itself, not just the seed that produced it.

    `seed` pins the order only for a fixed numpy. Kaggle updates its notebook image underneath us, and a
    session that resumed after such a bump would read a different permutation from the same seed -- some
    windows twice, others never -- while every other guard (weights, cursor arithmetic, fingerprint,
    batch size) still agreed. Hashing the full permutation is ~10 ms and makes that state impossible to
    enter quietly. Computed before the resume slice, so two sessions at different offsets agree on it.
    `base` lets a caller that already generated the permutation avoid generating it twice.
    """
    import hashlib
    if base is None:
        base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
    return hashlib.sha256(np.ascontiguousarray(base).tobytes()).hexdigest()[:16]


def make_dataset(store, seq_len, seed, start_sample=0, count=None, shuffle=True):
    """A map-style torch Dataset of (input_ids, labels, attention_mask) with an exact resume cursor.

    start_sample is the number of samples ALREADY consumed: the permutation is generated in full and
    then sliced, so the sample at index j of the resumed run is the same global sample it would have
    been in an uninterrupted run. That is the property Gate 3 has to demonstrate.
    """
    import torch

    n_samples = store.total_tokens // seq_len
    perm = None
    order_sha = "unshuffled"
    if shuffle:
        # Seed must be identical on every resume or the two runs read different data orders. Stored in
        # the checkpoint (cursor.json) and asserted by the trainer -- alongside perm_sha(), which catches
        # the case the seed alone cannot: the same seed on a different numpy.
        base = np.random.default_rng(np.random.PCG64(seed)).permutation(n_samples)
        order_sha = perm_sha(n_samples, seed, base=base)
        perm = base[start_sample:]
    else:
        perm = np.arange(start_sample, n_samples, dtype=np.int64)
    if count is not None:
        perm = perm[:count]

    class _DS(torch.utils.data.Dataset):
        def __len__(self):
            return int(len(perm))

        def __getitem__(self, j):
            s = int(perm[j])
            ids = store.tokens(s * seq_len, seq_len).astype(np.int64)
            x = torch.from_numpy(ids)
            return {"input_ids": x,
                    "labels": torch.cat((x[1:], x[:1])),
                    "attention_mask": torch.ones_like(x)}

        # Trainer calls remove_unused_columns on models whose forward signature it inspects; exposing
        # these names keeps it from dropping anything.
        column_names = ["input_ids", "labels", "attention_mask"]

    ds = _DS()
    # Instance attributes, not class attributes: `_DS.perm = perm` would let a second make_dataset() call
    # in the same process overwrite the first one's order.
    ds.perm = perm
    ds.order_sha = order_sha
    return ds


def cursor_dict(seq_len, samples_consumed, seed, dataset_files_sha=None, step=None, order_sha=None):
    """The whole of the data-position state. Written into every checkpoint, checked on every resume.

    `step` is recorded alongside `samples_consumed` so the two can be cross-checked on resume: they are
    redundant by construction (samples = step x sequences-per-step), and a cursor where they disagree is a
    corrupted or hand-edited file, not a position to continue from. `shuffle_perm_sha` is the identity of
    the visit order, which the seed alone does not pin (see perm_sha).
    """
    return {"seq_len": seq_len, "samples_consumed": samples_consumed,
            "tokens_consumed": samples_consumed * seq_len, "shuffle_seed": seed,
            "shuffle_perm_sha": order_sha, "step": step, "dataset_files_sha": dataset_files_sha}


def load_cursor(dirpath):
    p = os.path.join(dirpath, "cursor.json")
    if not os.path.exists(p):
        return None
    with open(p) as f:
        c = json.load(f)
    if not isinstance(c.get("samples_consumed"), int) or c["samples_consumed"] < 0:
        raise ValueError(f"{p}: samples_consumed missing or nonsense -- refusing a silent restart")
    return c


def save_cursor(dirpath, cur):
    os.makedirs(dirpath, exist_ok=True)
    with open(os.path.join(dirpath, "cursor.json"), "w") as f:
        json.dump(cur, f, indent=1, sort_keys=True)


def mix_fingerprint(root, manifest=None):
    """Cheap, stable identity for the exact data being trained on. Used by the trainer to refuse a resume
    onto a different corpus, so it has to be computed the same way in every session.

    Names alone (or names and sizes, as this function did until the E-031 review) are too weak: a rebuilt
    mix whose shards happen to land on the same sizes is a different corpus, and the whole point of the
    check is to catch a corpus swap between sessions. The manifest already certifies a sha256 per shard,
    so folding it in costs nothing and needs no file reads; the local size is included as well so a
    truncated file on disk cannot hide behind a manifest that still describes the good one.
    """
    import hashlib
    man = manifest or json.load(open(os.path.join(root, "manifest.json")))
    h = hashlib.sha256()
    for s in man["shards"]:
        p = os.path.join(root, s["file"])
        h.update(f"{s['file']}:{s.get('sha256')}:"
                 f"{os.path.getsize(p) if os.path.exists(p) else 'missing'}\n".encode())
    return h.hexdigest()[:16]