"""Tests for deterministic review sampling."""

from __future__ import annotations

from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq

from swe_data_processing.sampling import iter_review_records


def test_review_reads_only_provenance_shards(tmp_path: Path) -> None:
    selected_shard = tmp_path / "selected.parquet"
    unrelated_shard = tmp_path / "unrelated.parquet"
    pq.write_table(
        pa.Table.from_pylist(
            [
                {"trajectory_id": "wanted", "trajectory": "selected"},
                {"trajectory_id": "other", "trajectory": "same shard"},
            ]
        ),
        selected_shard,
    )
    pq.write_table(
        pa.Table.from_pylist(
            [{"trajectory_id": "unrelated", "trajectory": "must not be read"}]
        ),
        unrelated_shard,
    )
    selected = {"wanted": ["hard:FIVE_CONSECUTIVE_TOOL_FAILURES"]}
    decision = {
        "sample_id": "wanted",
        "source_parquet": str(selected_shard),
        "hard_reject_reasons": ["FIVE_CONSECUTIVE_TOOL_FAILURES"],
    }

    records = list(
        iter_review_records(
            tmp_path,
            selected,
            {"wanted": decision},
        )
    )

    assert [record["trajectory_id"] for record in records] == ["wanted"]
    assert records[0]["_qc_review"]["decision"] == decision