"""Arm A ingestion (argilla distilabel math preference pairs).
Run: uv run pytest tests/test_argilla_arm_a.py -m "not gpu" -v

Arm A is the off-policy, human/AI-rated preference arm. Its whole scientific
value depends on being EXACTLY the set that was measured at plan time, so the
counts are asserted, not logged.
"""

from __future__ import annotations

import json
import os

import pytest

from align.data import argilla, sources
from align.data.schema import PROMPT_TEMPLATE, stable_problem_id


def _rows():
    try:
        return argilla.load_argilla_rows()
    except (FileNotFoundError, RuntimeError) as e:
        pytest.skip(f"argilla snapshot unavailable: {e}")


def test_margin_filter_reproduces_measured_yield():
    """margin >=1.0 AND chosen >=6 keeps exactly 2,310 of 2,418 on the pinned
    dataset revision; a changed upstream revision fails the count assertion
    rather than silently training on different data (the drift incident
    class). Both directions: the pinned revision passes, a doctored row count
    raises."""
    rows = _rows()
    assert len(rows) == argilla.RAW_ROWS == 2418

    kept = argilla.apply_quality_filter(rows)
    assert len(kept) == argilla.FILTERED_ROWS == 2310

    # Firing direction: a drifted upstream (one row short) must abort.
    with pytest.raises(RuntimeError, match="revision"):
        argilla.assert_pinned_counts(rows[:-1], kept)


def test_contamination_zero_and_planted_problem_caught():
    """Contamination vs the eval hash set is 0/2,418 (silent direction); a
    deliberately planted eval problem IS caught (firing direction)."""
    rows = _rows()
    eval_hashes = {"deadbeefdeadbeef"}
    clean, n_contam = argilla.decontaminate(rows, eval_hashes)
    assert n_contam == 0
    assert len(clean) == len(rows)

    planted = stable_problem_id(rows[0]["instruction"])
    clean2, n2 = argilla.decontaminate(rows, eval_hashes | {planted})
    assert n2 >= 1
    assert len(clean2) == len(rows) - n2
    assert all(stable_problem_id(r["instruction"]) != planted for r in clean2)


def test_prompt_mapping_verbatim_no_answer_injection():
    """instruction maps into the prompt template's problem slot verbatim;
    completions are used verbatim -- no '# Answer' injection, no reformatting
    (format drift is a phenomenon to MEASURE at eval time, not to mask at
    build time)."""
    rows = _rows()
    pairs = argilla.to_pairs(argilla.apply_quality_filter(rows)[:32])
    for pair, row in zip(pairs, argilla.apply_quality_filter(rows)[:32]):
        assert pair.prompt == PROMPT_TEMPLATE.format(problem=row["instruction"].strip())
        assert row["instruction"].strip() in pair.prompt
        assert pair.chosen == row["chosen_response"]
        assert pair.rejected == row["rejected_response"]
        assert pair.problem_id == stable_problem_id(row["instruction"])

    # No injected answer marker anywhere: whatever the completions do or do
    # not contain is the measurement.
    injected = sum(
        1
        for p in pairs
        if "# Answer" in p.chosen
        and "# Answer"
        not in next(r["chosen_response"] for r in rows if r["instruction"] == p.instruction)
    )
    assert injected == 0


def test_length_bias_baseline_recorded(tmp_workdir):
    """Build receipt records chosen-longer fraction; equals 54.9% on the
    pinned revision (this number seeds the section-07 length-drift tripwire).
    A DPO run that drifts far above this is learning length, not correctness."""
    rows = _rows()
    kept = argilla.apply_quality_filter(rows)
    stats = argilla.length_bias_stats(kept)

    assert stats["chosen_longer_fraction"] == pytest.approx(0.549, abs=5e-4)
    assert stats["mean_chosen_chars"] == pytest.approx(1278.0, abs=1.0)
    assert stats["mean_rejected_chars"] == pytest.approx(1232.0, abs=1.0)

    receipt = argilla.write_arm_a_receipt(tmp_workdir, rows, kept, n_contaminated=0)
    on_disk = json.loads(open(os.path.join(tmp_workdir, "arm_a.receipt.json")).read())
    assert on_disk["chosen_longer_fraction"] == pytest.approx(0.549, abs=5e-4)
    assert on_disk["raw_rows"] == 2418
    assert on_disk["kept_rows"] == 2310
    assert on_disk["contaminated"] == 0
    assert on_disk["dataset_revision"] == sources.ARGILLA_PIN.revision
    assert receipt["descoped"] is False


def test_descope_marker_written_when_filter_starves_the_arm(tmp_workdir):
    """Pre-registered descope: a filter yielding < 1,500 pairs writes a
    DESCOPED marker instead of shards, so section-12 reports Arm A as
    descoped rather than reporting a null from an underpowered arm.
    Both directions: the real filter (2,310) writes no marker."""
    rows = _rows()

    starved = argilla.apply_quality_filter(rows, min_margin=1.0, min_chosen=9.5)
    assert len(starved) < argilla.DESCOPE_MIN_PAIRS
    marker = argilla.check_descope(tmp_workdir, starved)
    assert marker is not None
    assert os.path.exists(os.path.join(tmp_workdir, "arm_a.DESCOPED"))

    ok_dir = os.path.join(tmp_workdir, "ok")
    os.makedirs(ok_dir, exist_ok=True)
    assert argilla.check_descope(ok_dir, argilla.apply_quality_filter(rows)) is None
    assert not os.path.exists(os.path.join(ok_dir, "arm_a.DESCOPED"))