"""Regression tests for audit findings U09-1..U09-4 (align/data/store.py).
Run: uv run pytest tests/test_audit_g5_data_store_store.py -m "not gpu" -v
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest

from align.data.store import ShardDataset, ShardReader, ShardWriter, read_shard_receipt


def _add(w, ids, pid, **extra):
    w.add_row(ids, problem_id=pid, subject="Algebra", boundary_positions=[], **extra)


def _make_shard(d, name="s", rows=((1, 1, 1), (2, 2, 2, 2), (3, 3))):
    with ShardWriter(d, name) as w:
        for i, ids in enumerate(rows):
            _add(w, list(ids), f"p{i}")
    return w.receipt


# -- U09-1 --------------------------------------------------------------------


def test_u09_1_failed_add_row_leaves_no_partial_row(tmp_workdir):
    """A row rejected by schema/JSON validation must not leave its tokens in
    the blobs: the caller catches and continues (build.py's pattern), and the
    NEXT row's metadata would be paired with the PREVIOUS row's tokens."""
    w = ShardWriter(tmp_workdir, "s")
    _add(w, [1, 1, 1], "a", k=1)
    with pytest.raises(ValueError):
        _add(w, [2, 2, 2, 2], "b", k=1, extra_key=5)  # schema mismatch vs row 0
    with pytest.raises(TypeError):
        _add(w, [9, 9, 9, 9, 9], "c", k=np.int64(3))  # non-JSON extra
    _add(w, [3, 3], "d", k=2)
    w.finalize()
    rows = list(ShardReader(tmp_workdir, "s")[i] for i in range(2))
    assert [r.problem_id for r in rows] == ["a", "d"]
    assert rows[0].input_ids.tolist() == [1, 1, 1]
    assert rows[1].input_ids.tolist() == [3, 3]
    assert w.receipt.n_tokens == 5


def test_u09_1_non_json_extra_on_first_row_does_not_poison_schema(tmp_workdir):
    w = ShardWriter(tmp_workdir, "s")
    with pytest.raises(TypeError):
        _add(w, [1, 1], "a", bad=np.int64(3))
    _add(w, [2, 2], "b", good=1)  # schema is decided by the first ACCEPTED row
    w.finalize()
    r = ShardReader(tmp_workdir, "s")
    assert len(r) == 1 and r[0].problem_id == "b" and r[0].extras == {"good": 1}


def test_u09_1_ragged_field_failure_is_atomic_across_fields(tmp_workdir):
    w = ShardWriter(tmp_workdir, "s", {"input_ids": "int32", "rej": "int32"})
    w.add_row(
        [1, 1], problem_id="a", subject="x", boundary_positions=[], ragged={"rej": [5, 5]}
    )
    with pytest.raises(ValueError):
        w.add_row(
            [2, 2],
            problem_id="b",
            subject="x",
            boundary_positions=[],
            ragged={"rej": [[1, 2], [3, 4]]},  # 2-D: fails after input_ids passes
        )
    w.add_row(
        [3, 3, 3], problem_id="c", subject="x", boundary_positions=[], ragged={"rej": [6]}
    )
    w.finalize()
    r = ShardReader(tmp_workdir, "s")
    assert [r[i].problem_id for i in range(len(r))] == ["a", "c"]
    assert r[1].input_ids.tolist() == [3, 3, 3] and r[1].ragged["rej"].tolist() == [6]


# -- U09-2 --------------------------------------------------------------------


def test_u09_2_exception_in_with_block_leaves_no_receipt(tmp_workdir):
    with pytest.raises(KeyboardInterrupt):
        with ShardWriter(tmp_workdir, "s") as w:
            _add(w, [1, 1, 1], "a")
            raise KeyboardInterrupt
    assert w.receipt is None
    assert list(Path(tmp_workdir).glob("s.*")) == []
    with pytest.raises(FileNotFoundError):
        ShardDataset(tmp_workdir)


def test_u09_2_clean_exit_still_finalizes(tmp_workdir):
    r = _make_shard(tmp_workdir)
    assert r.n_rows == 3 and (Path(tmp_workdir) / "s.receipt.json").exists()


def test_u09_2_finalized_shard_survives_a_later_exception(tmp_workdir):
    with pytest.raises(RuntimeError):
        with ShardWriter(tmp_workdir, "s") as w:
            _add(w, [1, 1], "a")
            w.finalize()
            raise RuntimeError("after finalize")
    assert (Path(tmp_workdir) / "s.receipt.json").exists()


def test_u09_2_poisoned_writer_refuses_to_finalize(tmp_workdir):
    """A write that failed AFTER the first byte hit disk cannot be undone."""
    w = ShardWriter(tmp_workdir, "s")
    _add(w, [1, 1], "a")
    w._handles["input_ids"].close()  # next write raises ValueError mid-row
    with pytest.raises(ValueError):
        _add(w, [2, 2], "b")
    with pytest.raises(RuntimeError, match="poison|abort|partial"):
        _add(w, [3, 3], "c")
    with pytest.raises(RuntimeError, match="poison|abort|partial"):
        w.finalize()


# -- U09-3 --------------------------------------------------------------------


def test_u09_3_tampered_idx_is_detected(tmp_workdir):
    _make_shard(tmp_workdir)
    idx = Path(tmp_workdir) / "s.input_ids.idx.npy"
    np.save(idx, np.array([0, 3, 7], dtype=np.int64))  # missing last offset
    with pytest.raises((RuntimeError, ValueError)):
        r = ShardReader(tmp_workdir, "s")
        r[0]


def test_u09_3_trailing_bytes_in_bin_are_detected(tmp_workdir):
    _make_shard(tmp_workdir)
    with open(Path(tmp_workdir) / "s.input_ids.bin", "ab") as f:
        f.write(b"\0" * 8)
    with pytest.raises((RuntimeError, ValueError)):
        r = ShardReader(tmp_workdir, "s")
        r[0]


def test_u09_3_shifted_row_boundaries_change_the_hash(tmp_workdir):
    """Same length, same last offset, different row boundaries: only a hash
    that covers the idx files notices."""
    _make_shard(tmp_workdir)
    r = ShardReader(tmp_workdir, "s")
    r.verify()  # untampered shard verifies
    idx = Path(tmp_workdir) / "s.input_ids.idx.npy"
    np.save(idx, np.array([0, 2, 9, 9], dtype=np.int64))
    with pytest.raises(RuntimeError, match="hash"):
        ShardReader(tmp_workdir, "s").verify()


def test_u09_3_legacy_v1_receipt_still_verifies(tmp_workdir):
    """Shards built before the idx files were hashed carry format_version 1 and
    must keep verifying against the file set they were hashed over."""
    from align.data.store import _hash_files

    _make_shard(tmp_workdir)
    d = Path(tmp_workdir)
    legacy = _hash_files(
        [d / f"s.{f}.bin" for f in ("boundary_positions", "input_ids")] + [d / "s.extras.jsonl"]
    )
    rj = json.loads((d / "s.receipt.json").read_text())
    rj["content_hash"], rj["format_version"] = legacy, 1
    (d / "s.receipt.json").write_text(json.dumps(rj))
    ShardReader(tmp_workdir, "s").verify()


# -- U09-4 --------------------------------------------------------------------


def test_u09_4_rewriting_a_shard_drops_its_stale_receipt_immediately(tmp_workdir):
    """Between `open(wb)` truncating the blobs and finalize(), the OLD receipt
    must not be left describing bytes that no longer exist."""
    _make_shard(tmp_workdir)
    w = ShardWriter(tmp_workdir, "s")
    assert not (Path(tmp_workdir) / "s.receipt.json").exists()
    w.abort()