"""Ragged shard store + receipts.
Run: uv run pytest tests/test_data_store.py -m "not gpu" -v

The store is the only place training data enters the run, so it is where the
"missing value silently replaced by a plausible one" failure class has to be
made impossible: `problem_id`, `subject` and the per-row boundary receipt are
MANDATORY keyword arguments with no defaults, validated at write time.
"""

from __future__ import annotations

import json
import os

import numpy as np
import pytest

from align.data import tokenization as tk
from align.data.store import (
    ShardDataset,
    ShardReader,
    ShardWriter,
    read_shard_receipt,
)

NL = 781  # test-side constant; the PACKAGE derives it, never hardcodes it


def _spec() -> tk.BoundarySpec:
    return tk.BoundarySpec(newline_id=NL, tokenizer_hash="test", vocab_absorbers=0)


def _row_ids(n_steps: int) -> list[int]:
    """A synthetic token row with `n_steps - 1` newline-bigram boundaries."""
    out: list[int] = []
    for k in range(n_steps):
        if k:
            out += [NL, NL]
        out += [100 + k, 200 + k, 300 + k]
    return out


def test_shard_metadata_carries_problem_id_and_subject(tmp_workdir):
    """Every row written to a shard round-trips problem_id and subject via
    metadata extras.
    Incident: the predecessor added these post-hoc -- subject was `null` on all
    10,606 RL prompts and Stage 4 would have aborted AFTER loading a 7B model.
    Here they are build-time mandatory and travel with the row."""
    spec = _spec()
    with ShardWriter(tmp_workdir, "sft-0000") as w:
        for i in range(5):
            ids = _row_ids(i + 2)
            w.add_row(
                ids,
                problem_id=f"pid{i:04d}",
                subject=["Algebra", "Geometry"][i % 2],
                boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
            )
    receipt = w.receipt
    assert receipt.n_rows == 5

    r = ShardReader(tmp_workdir, "sft-0000")
    assert len(r) == 5
    for i in range(5):
        row = r[i]
        assert row.problem_id == f"pid{i:04d}"
        assert row.subject == ["Algebra", "Geometry"][i % 2]
        assert row.input_ids.tolist() == _row_ids(i + 2)

    # The receipt is on disk and reloads identically -- the manifest quotes it.
    assert read_shard_receipt(tmp_workdir, "sft-0000") == receipt


def test_receipts_roundtrip_against_fresh_scan(tmp_workdir):
    """Boundary positions stored in shard extras match a fresh integer scan of
    the stored input_ids for EVERY row of a small built shard (uses the
    section-02 scan; no decode anywhere).
    Incident: a receipt that drifts from the ids it describes is undetectable
    downstream -- interp would patch the wrong positions and report a real
    number about the wrong tokens."""
    spec = _spec()
    with ShardWriter(tmp_workdir, "sft-0000") as w:
        for i in range(8):
            ids = _row_ids(i + 1)
            w.add_row(
                ids,
                problem_id=f"p{i}",
                subject="Algebra",
                boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
            )

    ds = ShardDataset(tmp_workdir)
    assert len(ds) == 8
    checked = 0
    for i in range(len(ds)):
        row = ds[i]
        fresh = tk.find_step_boundaries(row.input_ids.tolist(), spec.newline_id)
        assert row.boundary_positions.tolist() == fresh
        checked += 1
    assert checked == 8


def test_row_missing_mandatory_extras_raises(tmp_workdir):
    """Writing a row without problem_id/subject/boundary receipts raises; the
    store never backfills a plausible default (the silent-substitution failure
    class). Both directions: a complete row writes silently."""
    spec = _spec()
    ids = _row_ids(3)
    good = tk.find_step_boundaries(ids, spec.newline_id)

    with ShardWriter(tmp_workdir, "s0") as w:
        with pytest.raises(TypeError):
            w.add_row(ids, subject="Algebra", boundary_positions=good)
        with pytest.raises(TypeError):
            w.add_row(ids, problem_id="p", boundary_positions=good)
        with pytest.raises(TypeError):
            w.add_row(ids, problem_id="p", subject="Algebra")
        for bad in ("", None):
            with pytest.raises(ValueError):
                w.add_row(ids, problem_id=bad, subject="Algebra", boundary_positions=good)
            with pytest.raises(ValueError):
                w.add_row(ids, problem_id="p", subject=bad, boundary_positions=good)
        # A receipt pointing outside the row is a corrupt receipt, not a row
        # to be quietly clipped.
        with pytest.raises(ValueError):
            w.add_row(ids, problem_id="p", subject="Algebra", boundary_positions=[len(ids) + 5])
        with pytest.raises(ValueError):
            w.add_row(
                ids, problem_id="p", subject="Algebra", boundary_positions=[good[-1], good[0]]
            )  # not increasing
        # Silent direction: the complete row goes through.
        w.add_row(ids, problem_id="p", subject="Algebra", boundary_positions=good)

    assert w.receipt.n_rows == 1


def test_content_hash_detects_a_flipped_token(tmp_workdir):
    """The shard content hash is what the build manifest pins, so it must
    change when a single stored token changes. Both directions: an identical
    rebuild reproduces the hash exactly."""
    spec = _spec()

    def build(dirname: str, tweak: int) -> str:
        d = os.path.join(tmp_workdir, dirname)
        with ShardWriter(d, "s") as w:
            ids = _row_ids(3)
            ids[0] += tweak
            w.add_row(
                ids,
                problem_id="p",
                subject="Algebra",
                boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
            )
        return w.receipt.content_hash

    a = build("a", 0)
    b = build("b", 0)
    c = build("c", 1)
    assert a == b  # deterministic rebuild
    assert a != c  # one flipped token changes the hash


def test_extras_sidecar_is_one_json_object_per_row(tmp_workdir):
    """The extras sidecar is line-delimited so a shard streams; a row count
    mismatch between the binary blob and the sidecar raises at open time
    rather than mis-pairing metadata with tokens."""
    spec = _spec()
    with ShardWriter(tmp_workdir, "s") as w:
        for i in range(3):
            ids = _row_ids(2)
            w.add_row(
                ids,
                problem_id=f"p{i}",
                subject="Algebra",
                boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
            )

    path = os.path.join(tmp_workdir, "s.extras.jsonl")
    lines = [json.loads(ln) for ln in open(path) if ln.strip()]
    assert len(lines) == 3
    assert {"problem_id", "subject", "boundary_positions"} <= set(lines[0])

    with open(path, "a") as f:
        f.write(
            json.dumps({"problem_id": "ghost", "subject": "Algebra", "boundary_positions": []})
            + "\n"
        )
    with pytest.raises(RuntimeError, match="row count"):
        ShardReader(tmp_workdir, "s")


def test_lengths_read_from_index_only(tmp_workdir):
    """Length-grouped batching scans every row length each epoch; that scan
    must not touch the data pages (it would page in the whole corpus on a box
    with no swap)."""
    spec = _spec()
    with ShardWriter(tmp_workdir, "s") as w:
        for i in range(6):
            ids = _row_ids(i + 1)
            w.add_row(
                ids,
                problem_id=f"p{i}",
                subject="Algebra",
                boundary_positions=tk.find_step_boundaries(ids, spec.newline_id),
            )

    ds = ShardDataset(tmp_workdir)
    lengths = ds.lengths()
    assert lengths.tolist() == [len(_row_ids(i + 1)) for i in range(6)]
    assert isinstance(lengths, np.ndarray)