"""Regression tests for audit findings U09-4..U09-7 (align/data/build.py,
align/data/prm800k_pairs.py).
Run: uv run pytest tests/test_audit_g5_data_store_build.py -m "not gpu" -v
"""

from __future__ import annotations

import json
from pathlib import Path

import pytest

from align.data import build as bd
from align.data import prm800k_pairs as pb
from align.data.build import _PREF_FIELDS, _ShardSet, _held_out_ids
from align.data.store import ShardDataset


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


def _fill(set_, pids):
    for pid in pids:
        set_.add(input_ids=[1, 2, 3], problem_id=pid, subject="s", boundary_positions=[])
    return set_.finalize()


def test_u09_4_rebuild_into_reused_dir_does_not_leave_stale_shards(tmp_workdir):
    old = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3)
    _fill(old, ["o0", "o1", "o2", "o3", "o4"])  # shards 0000 + 0001
    new = _ShardSet(tmp_workdir, "sft-train", rows_per_shard=3)
    _fill(new, ["n0"])  # only 0000 rewritten
    ids = [r.problem_id for r in ShardDataset(tmp_workdir).iter_rows()]
    assert ids == ["n0"]


def test_u09_4_purge_is_scoped_to_its_own_prefix(tmp_workdir):
    other = _ShardSet(tmp_workdir, "sft-held")
    _fill(other, ["h0"])
    _fill(_ShardSet(tmp_workdir, "sft-train"), ["t0"])
    assert [r.problem_id for r in ShardDataset(tmp_workdir, names=["sft-held-0000"]).iter_rows()] == ["h0"]


# -- U09-5 --------------------------------------------------------------------


def _ids(n):
    return [f"problem-{i:06d}" for i in range(n)]


def test_u09_5_held_out_assignment_is_independent_of_the_id_set():
    ids = _ids(4000)
    held_full = _held_out_ids(ids, 0.05, 0)
    # Arm-B-like subset: every id's membership must be the same as in the
    # full set, otherwise Arm B's held-out set leaks SFT-train problems.
    subset = ids[::3]
    held_sub = _held_out_ids(subset, 0.05, 0)
    assert held_sub == held_full & set(subset)


def test_u09_5_adding_one_id_changes_at_most_that_id():
    ids = _ids(4000)
    base = _held_out_ids(ids, 0.05, 0)
    more = _held_out_ids(ids + ["problem-new"], 0.05, 0)
    assert base - more == set()
    assert more - base <= {"problem-new"}


def test_u09_5_fraction_is_close_to_requested_and_nonempty_for_tiny_inputs():
    held = _held_out_ids(_ids(20000), 0.05, 0)
    assert 0.04 * 20000 < len(held) < 0.06 * 20000
    assert len(_held_out_ids(["only-one"], 0.05, 0)) == 1
    assert len(_held_out_ids(_ids(5), 0.05, 0)) >= 1
    assert _held_out_ids([], 0.05, 0) == set()


def test_u09_5_seed_changes_the_split():
    assert _held_out_ids(_ids(4000), 0.05, 0) != _held_out_ids(_ids(4000), 0.05, 1)


def test_u09_5_rule_is_recorded_in_pins():
    assert bd.PinnedYields().held_out_rule == "hash-v1"


# -- U09-6 --------------------------------------------------------------------


def _pair(pid="p0", step=0, chosen="good", rejected="bad", conf=1.0):
    return pb.PrefPair(
        problem="Q", problem_id=pid, step_index=step, prefix_steps=(),
        chosen_step=chosen, rejected_step=rejected, confidence=conf,
    )


def test_u09_6_curate_never_emits_duplicate_pair_ids():
    dup = [_pair(), _pair(), _pair(rejected="worse")]
    selected, remainder = pb.curate(dup, set(), {"p0": "Algebra"}, target=1)
    ids = [p.pair_id for p in selected] + [p.pair_id for p in remainder]
    assert len(ids) == len(set(ids)) == 2
    assert not ({p.pair_id for p in selected} & {p.pair_id for p in remainder})


def test_u09_6_curate_dedupe_is_order_independent():
    a, b = _pair(), _pair(rejected="worse")
    s1, r1 = pb.curate([a, a, b], set(), {"p0": "Algebra"}, target=1)
    s2, r2 = pb.curate([b, a, a], set(), {"p0": "Algebra"}, target=1)
    assert [p.pair_id for p in s1] == [p.pair_id for p in s2]
    assert [p.pair_id for p in r1] == [p.pair_id for p in r2]


# -- U09-7 --------------------------------------------------------------------


class _FakeSpec:
    newline_id = 1
    vocab_absorbers = 0
    tokenizer_hash = "t"

    def to_json(self, path):
        Path(path).write_text("{}")


def test_u09_7_aborted_rebuild_invalidates_the_previous_manifest(tmp_workdir, monkeypatch):
    out = Path(tmp_workdir) / "build"
    out.mkdir()
    (out / "manifest.json").write_text(json.dumps({"pins_verified": True, "build_version": 2}))
    monkeypatch.setattr(bd, "derive_boundary_spec", lambda tok: _FakeSpec())
    monkeypatch.setattr(bd, "math_subject_index", lambda src: {})

    def boom(*a, **k):
        raise RuntimeError("strict pin failure")

    monkeypatch.setattr(bd, "build_eval_stage", boom)
    with pytest.raises(RuntimeError, match="strict pin"):
        bd.build_all(object(), out, object(), bd.PinnedYields())
    assert not (out / "manifest.json").exists()


def test_u09_7_atomic_write_never_exposes_a_partial_file(tmp_workdir, monkeypatch):
    from align.data import store

    target = Path(tmp_workdir) / "m.json"
    target.write_text("OLD")

    def fail_replace(src, dst):
        raise OSError("killed before rename")

    monkeypatch.setattr(store.os, "replace", fail_replace)
    with pytest.raises(OSError):
        store.atomic_write_text(target, "NEW")
    assert target.read_text() == "OLD"
    assert [p.name for p in Path(tmp_workdir).iterdir()] == ["m.json"]
    monkeypatch.undo()
    store.atomic_write_text(target, "NEW")
    assert target.read_text() == "NEW"