"""build_data.py end-to-end on miniature inputs.
Run: uv run pytest tests/test_build_data_e2e.py -m "not gpu" -v

The build has ONE legal order -- eval sets, then the decontamination hash set,
then every training corpus excluded against it -- because any other order
produces a corpus that was never checked against the thing it will be
evaluated on, and nothing downstream can tell the difference.
"""

from __future__ import annotations

import json
import math
import os

import pytest

from align.data import build as bd
from align.data.sources import DataSources

# ---------------------------------------------------------------------------
# Miniature source fixtures: same file layout as the real snapshots, ~20 rows.
# ---------------------------------------------------------------------------


def _write_rows(path: str, rows: list[dict]) -> None:
    """JSONL fixture (sources.iter_local_rows reads it like a parquet; the
    suite stays pyarrow-free)."""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as f:
        for r in rows:
            f.write(json.dumps(r) + "\n")


@pytest.fixture
def mini_sources(tmp_workdir) -> DataSources:
    root = os.path.join(tmp_workdir, "src")
    os.makedirs(root, exist_ok=True)

    # --- eval sets -------------------------------------------------------
    m500 = os.path.join(root, "math500", "test.jsonl")
    os.makedirs(os.path.dirname(m500), exist_ok=True)
    with open(m500, "w") as f:
        for i in range(6):
            f.write(
                json.dumps(
                    {
                        "problem": f"Eval math problem {i}. What is ${i}+1$?",
                        "answer": str(i + 1),
                        "subject": "Algebra",
                        "level": 3,
                        "solution": "trivial",
                        "unique_id": f"test/algebra/{i}.json",
                    }
                )
                + "\n"
            )

    gsm = os.path.join(root, "gsm8k", "test-00000-of-00001.jsonl")
    _write_rows(
        gsm,
        [
            {"question": f"Eval gsm problem {i}?", "answer": f"reasoning\n#### {i}"}
            for i in range(4)
        ],
    )

    # --- MATH corpus for the subject join --------------------------------
    for subject, n in (("algebra", 6), ("geometry", 6)):
        _write_rows(
            os.path.join(root, "hendrycks", subject, "train-00000-of-00001.jsonl"),
            [
                {
                    "problem": f"Train problem {subject} {i}.",
                    "solution": "s",
                    "type": subject.capitalize(),
                    "level": "Level 3",
                }
                for i in range(n)
            ],
        )

    # --- PRM800K annotation records --------------------------------------
    def node(comps, chosen=None):
        return {
            "completions": [{"text": t, "rating": r, "flagged": False} for t, r in comps],
            "chosen_completion": chosen,
            "human_completion": None,
        }

    recs = []
    for subject in ("algebra", "geometry"):
        for i in range(6):
            problem = f"Train problem {subject} {i}."
            # A clean solution (SFT-eligible) ...
            recs.append(
                {
                    "question": {"problem": problem, "ground_truth_answer": str(i)},
                    "label": {
                        "finish_reason": "solution",
                        "steps": [
                            node([(f"First we consider case {i}.", 1)], chosen=0),
                            node([(f"# Answer\n\n{i}", 1)], chosen=0),
                        ],
                    },
                    "is_quality_control_question": False,
                }
            )
            # ... and an annotation with a rated sibling pair (Arm B site).
            recs.append(
                {
                    "question": {"problem": problem, "ground_truth_answer": str(i)},
                    "label": {
                        "finish_reason": "found_error",
                        "steps": [
                            node([(f"First we consider case {i}.", 1)], chosen=0),
                            node(
                                [
                                    (f"Then the answer is {i}.", 1),
                                    (f"Then the answer is {i + 99}.", -1),
                                    ("This restates the problem.", 0),
                                ],
                                chosen=0,
                            ),
                        ],
                    },
                    "is_quality_control_question": False,
                }
            )
    # One record that is contaminated: it is a MATH-500 eval problem.
    recs.append(
        {
            "question": {
                "problem": "Eval math problem 0. What is $0+1$?",
                "ground_truth_answer": "1",
            },
            "label": {
                "finish_reason": "solution",
                "steps": [
                    node([("Leaked step.", 1)], chosen=0),
                    node([("# Answer\n\n1", 1)], chosen=0),
                ],
            },
            "is_quality_control_question": False,
        }
    )

    prm_dir = os.path.join(root, "prm800k")
    os.makedirs(prm_dir, exist_ok=True)
    with open(os.path.join(prm_dir, "phase1_train.jsonl"), "w") as f:
        for r in recs[:4]:
            f.write(json.dumps(r) + "\n")
    with open(os.path.join(prm_dir, "phase2_train.jsonl"), "w") as f:
        for r in recs[4:]:
            f.write(json.dumps(r) + "\n")

    # --- argilla preference rows -----------------------------------------
    arg = os.path.join(root, "argilla", "train.jsonl")
    _write_rows(
        arg,
        [
            {
                "metadata": "{}",
                "instruction": f"Argilla problem {i}?",
                "chosen_response": "A longer chosen response with detail. " * (2 + i % 2),
                "chosen_rating": 9.0 if i % 4 else 5.0,
                "rejected_response": "Short rejected response.",
                "rejected_rating": 7.0,
            }
            for i in range(8)
        ],
    )

    return DataSources(
        prm800k_dir=prm_dir,
        math500_file=m500,
        gsm8k_test_file=gsm,
        hendrycks_math_dir=os.path.join(root, "hendrycks"),
        argilla_file=arg,
        revisions={
            "prm800k": "mini",
            "math500": "mini",
            "gsm8k": "mini",
            "hendrycks_math": "mini",
            "argilla": "mini",
        },
    )


@pytest.fixture
def mini_pins() -> bd.PinnedYields:
    """The miniature build's own pinned counts. Production defaults are the
    measured real-corpus numbers; a test that reused them would either be
    skipped or would teach us to relax the assertion."""
    return bd.PinnedYields(
        math500=6,
        gsm8k=4,
        decontamination=10,
        sft_rows=12,
        argilla_raw=8,
        argilla_filtered=6,
        arm_b_sites=12,
        arm_b_capped=12,
        arm_b_problems=12,
        arm_b_curated=8,
        descope_min_pairs=1,
        expected_chosen_longer=None,
        held_out_frac=0.2,
    )


def test_strict_build_order_and_receipts(mini_sources, mini_pins, tmp_workdir, v03_tokenizer):
    """Eval sets build FIRST, then the decontamination hash set, then training
    corpora excluded against it, then shards. Every artifact has a receipt;
    the final manifest lists all hashes + counts. Attempting to build a
    training corpus before the eval hash set exists raises."""
    out = os.path.join(tmp_workdir, "data")

    # Firing direction: the SFT stage refuses to run without its upstream.
    with pytest.raises(FileNotFoundError, match="eval"):
        bd.build_sft_corpus(mini_sources, out, v03_tokenizer, mini_pins)

    manifest = bd.build_all(mini_sources, out, v03_tokenizer, pins=mini_pins)

    assert manifest["stages"] == ["eval_sets", "decontamination", "sft", "arm_a", "arm_b"]
    assert manifest["counts"]["math500"] == 6
    assert manifest["counts"]["gsm8k_test"] == 4
    assert manifest["counts"]["decontamination_hashes"] == 10
    # The planted contaminated record is dropped, loudly and countably.
    assert manifest["counts"]["sft_rows"] == 12
    assert manifest["counts"]["sft_contaminated"] == 1
    assert manifest["counts"]["sft_rejected_boundary"] == 0

    # Stamped from a constant, never by hand: section-11's gates require
    # build_version >= 2 and RAISE on a missing field, so an unstamped build is
    # refused rather than assumed current. 1 = the doubled-marker builds.
    assert manifest["build_version"] == bd.BUILD_VERSION >= 2

    assert manifest["boundary_spec"]["vocab_absorbers"] == 0
    assert manifest["tokenizer_hash"]
    for key in ("math500", "gsm8k_test", "sft-train", "arm_a-train", "arm_b-train"):
        assert manifest["hashes"][key], f"missing content hash for {key}"

    # The corpus outcome-agreement statistic is RECORDED, with its decision
    # attached -- it is a measured property of the data, not a gate, and the
    # manifest is where it stays visible across rebuilds.
    # Arm B carries its OWN length baseline: its pairs share a step prefix, so
    # Arm A's number would arm the tripwire against a different distribution.
    lb = manifest["length_bias_baseline"]
    assert not math.isnan(lb["arm_b_chosen_longer_fraction"])
    assert not math.isnan(lb["arm_b_chosen_longer_fraction_step_basis"])
    # Arm A's pinned baseline is CHARACTER-based and must not be overwritten by
    # the token-based measurement the shard writer computes.
    assert lb["arm_a_chosen_longer_fraction"] == pytest.approx(
        manifest["counts"].get("arm_a_chosen_longer_fraction", lb["arm_a_chosen_longer_fraction"])
    )

    agreement = manifest["corpus_outcome_agreement"]
    assert (
        agreement["n"]
        == agreement["agree"]
        + agreement["equivalent_unnormalized"]
        + agreement["genuinely_divergent"]
    )
    assert agreement["rows_kept"] is True
    assert "anchor" in agreement["note"]

    on_disk = json.loads(open(os.path.join(out, "manifest.json")).read())
    assert on_disk == manifest


def test_split_by_problem_never_by_record(mini_sources, mini_pins, tmp_workdir, v03_tokenizer):
    """No problem_id appears on both sides of any split (SFT vs held-out, pair
    train vs held-out preference-accuracy set).
    Incident: PRM800K holds many annotation records per problem, so a
    record-level split reports an inflated validation number that looks like
    generalization."""
    out = os.path.join(tmp_workdir, "data")
    bd.build_all(mini_sources, out, v03_tokenizer, pins=mini_pins)

    for corpus in ("sft", "arm_b"):
        train = bd.problem_ids_in(os.path.join(out, "shards", corpus, "train"))
        held = bd.problem_ids_in(os.path.join(out, "shards", corpus, "held_out"))
        assert train, f"{corpus} train split is empty"
        assert held, f"{corpus} held-out split is empty"
        assert not (train & held), f"{corpus} leaks problems across the split"


def test_all_outputs_under_apps(tmp_workdir):
    """Every writable path statvfs-resolves onto the /apps device (uses the
    section-01 guard; tested here in the firing direction with a doctored
    output dir).
    Incident: / has ~13GB free and /tmp is a noexec tmpfs on a DIFFERENT
    device, so 'not the root device' is not enough -- it must be /apps."""
    ok = os.path.join(tmp_workdir, "data")
    bd.assert_build_paths(bd.output_paths(ok))  # silent direction

    doctored = "/tmp/nmo-should-never-write-here"
    try:
        with pytest.raises(RuntimeError, match="/apps"):
            bd.assert_build_paths(bd.output_paths(doctored))
    finally:
        # The guard creates each path before stat'ing it, so the firing
        # direction leaves a directory on the very filesystem this repo is
        # forbidden to write to. Remove it.
        import shutil

        shutil.rmtree(doctored, ignore_errors=True)