"""End-to-end tests for the two-stage ablation pipeline.

The agent is simulated: a world where some sections genuinely help, some
genuinely hurt, and most do nothing. What is under test is whether the
pipeline finds them and — more important — whether it ever recommends deleting
one that matters.
"""

from __future__ import annotations

import io

import numpy as np
import pytest
from rich.console import Console

from ablate.ablation import ablatable, confirm, screen, select_candidates
from ablate.ablation.pipeline import AblationResult
from ablate.config.parse import split_sections
from ablate.config.variant import ConfigVariant
from ablate.errors import AblateError
from ablate.record import AssertionRecord, TelemetryRecord, TrialRecord
from ablate.report.ablation import render_ablation
from ablate.stats import Verdict

TASKS = [f"task-{i:02d}" for i in range(12)]

BODY = (
    "Guidance about how to write code in this repository, stated at enough "
    "length that it clears the small-section merge threshold comfortably. "
)


def _config(n: int = 6, pinned: set[int] | None = None) -> str:
    """Sections of deliberately different sizes.

    Equal-sized sections would make "the largest deletion candidate"
    ambiguous, and a tie broken by sort order is not a behaviour worth
    pinning.
    """
    pinned = pinned or set()
    parts = []
    for i in range(n):
        marker = "<!-- ablate:keep -->\n" if i in pinned else ""
        parts.append(f"## Section {i}\n\n{marker}{BODY * (i + 1)}\n")
    return "\n".join(parts)


class SimulatedAgent:
    """Pass probability is a base difficulty plus the effect of each present section."""

    def __init__(self, sections, effects: dict[str, float], seed: int = 0, tokens_per_section=40):
        self.sections = list(sections)
        self.effects = effects
        self.rng = np.random.default_rng(seed)
        self.difficulty = {t: d for t, d in zip(TASKS, self.rng.uniform(0.25, 0.75, len(TASKS)), strict=True)}
        self.tokens_per_section = tokens_per_section
        self.calls: list[tuple[int, int, str]] = []

    def _present(self, variant: ConfigVariant) -> list[str]:
        return [s.id for s in self.sections if s.title in (variant.text or "")]

    def __call__(self, variants, trials, stage):
        self.calls.append((len(variants), trials, stage))
        records = []
        for variant in variants:
            present = self._present(variant)
            lift = sum(self.effects.get(sid, 0.0) for sid in present)
            for task in TASKS:
                for replicate in range(1, trials + 1):
                    p = float(np.clip(self.difficulty[task] + lift, 0.02, 0.98))
                    passed = bool(self.rng.random() < p)
                    turns = 6
                    tokens = (200 + self.tokens_per_section * len(present)) * turns
                    records.append(
                        TrialRecord(
                            trial_id=f"{task}-{variant.id}-r{replicate}",
                            run_id="run1",
                            task_id=task,
                            variant_id=variant.id,
                            variant_hash="sha256:x",
                            replicate=replicate,
                            agent="fake",
                            model="fake-model",
                            started_at="2026-07-27T00:00:00Z",
                            duration_s=1.0,
                            telemetry=TelemetryRecord(
                                input_tokens=tokens, turns=turns, cost_usd=0.01
                            ),
                            assertions=[
                                AssertionRecord(id="a", type="exit_code", passed=passed)
                            ],
                            passed=passed,
                        )
                    )
        return records


def _sections(n: int = 6, pinned: set[int] | None = None):
    return split_sections(_config(n, pinned))


#: Each bootstrap resample re-runs leave-one-task-out CV to pick its own λ, so
#: the production default of 2,000 would put this file in the minutes. These
#: tests check structure and ranking, not interval precision — the calibration
#: figures live in test_regression.py, which pays for the resamples it needs.
TEST_RESAMPLES = 60


# --- eligibility ------------------------------------------------------------


def test_pinned_sections_are_not_ablatable():
    sections = _sections(4, pinned={1})
    assert len(sections) == 4
    assert len(ablatable(sections)) == 3


def test_screening_refuses_when_everything_is_pinned():
    sections = _sections(3, pinned={0, 1, 2})
    with pytest.raises(AblateError, match="nothing to ablate"):
        screen(sections, SimulatedAgent(sections, {}), model="fake", seed=1, n_resamples=TEST_RESAMPLES)


# --- stage 1 ----------------------------------------------------------------


def test_screening_covers_every_ablatable_section():
    sections = _sections(6)
    result = screen(sections, SimulatedAgent(sections, {}, seed=1), model="fake", seed=1, n_resamples=TEST_RESAMPLES)
    assert len(result.rows) == 6
    assert {r.section.id for r in result.rows} == {s.id for s in sections}


def test_screening_runs_one_trial_per_variant_task_cell():
    sections = _sections(5)
    agent = SimulatedAgent(sections, {}, seed=1)
    result = screen(sections, agent, model="fake", seed=1, n_resamples=TEST_RESAMPLES)
    n_variants, trials, stage = agent.calls[0]
    assert (n_variants, trials, stage) == (2 * 5 + 8, 1, "screening")
    assert result.n_records == n_variants * len(TASKS)


def test_screening_ranks_the_section_that_matters_first():
    sections = _sections(6)
    effects = {"s02": 0.35}
    result = screen(sections, SimulatedAgent(sections, effects, seed=2), model="fake", seed=2, n_resamples=TEST_RESAMPLES)
    assert result.ranked[0].section.id == "s02"


def test_screening_recovers_the_empty_config_floor():
    """The all-dropped variant answers "what if I delete the whole thing?",
    which is the question most users actually have."""
    sections = _sections(5)
    effects = {s.id: 0.06 for s in sections}  # every section helps a little
    result = screen(sections, SimulatedAgent(sections, effects, seed=4), model="fake", seed=4, n_resamples=TEST_RESAMPLES)
    assert result.floor_pass_rate < result.reference_pass_rate
    assert result.floor_delta > 0


def test_pinned_sections_appear_in_every_screening_variant():
    sections = _sections(5, pinned={0})
    agent = SimulatedAgent(sections, {}, seed=1)
    screen(sections, agent, model="fake", seed=1, n_resamples=TEST_RESAMPLES)
    assert all(s.id != sections[0].id for s in [])  # sanity: pinned excluded from design

    # Every rendered variant, including the "floor", still contains the pin.
    from ablate.ablation.design import build_design
    from ablate.config.parse import render_variant

    design = build_design([s.id for s in ablatable(sections)], seed=1)
    for mask in design.masks:
        chosen = dict(zip((s.id for s in ablatable(sections)), mask, strict=True))
        rendered = render_variant(sections, [chosen.get(s.id, True) for s in sections])
        assert sections[0].title in rendered


def test_screening_raises_when_every_trial_errored():
    sections = _sections(4)

    def broken(variants, trials, stage):
        return []

    with pytest.raises(AblateError, match="no usable trials"):
        screen(sections, broken, model="fake", seed=1, n_resamples=TEST_RESAMPLES)


# --- candidate selection ----------------------------------------------------


def test_candidates_include_the_strongest_mover():
    sections = _sections(6)
    result = screen(sections, SimulatedAgent(sections, {"s02": 0.35}, seed=2), model="fake", seed=2, n_resamples=TEST_RESAMPLES)
    assert "s02" in {s.id for s in select_candidates(result, limit=3)}


def test_candidates_include_the_largest_deletion_candidate():
    """Big sections whose effect is consistent with zero are the ones a user
    would actually act on."""
    sections = _sections(6)
    result = screen(sections, SimulatedAgent(sections, {}, seed=3), model="fake", seed=3, n_resamples=TEST_RESAMPLES)
    candidates = select_candidates(result, limit=3)
    deletion_rows = sorted(
        (r for r in result.rows if r.is_deletion_candidate), key=lambda r: -r.tokens
    )
    assert deletion_rows[0].section.id in {s.id for s in candidates}
    # ...and it is genuinely the biggest, not a tie broken by ordering.
    assert deletion_rows[0].tokens > deletion_rows[-1].tokens


def test_candidate_count_is_capped():
    sections = _sections(8)
    result = screen(sections, SimulatedAgent(sections, {}, seed=1), model="fake", seed=1, n_resamples=TEST_RESAMPLES)
    assert len(select_candidates(result, limit=2)) == 2
    assert len(select_candidates(result, limit=0)) == 0


def test_candidates_are_distinct():
    sections = _sections(6)
    result = screen(sections, SimulatedAgent(sections, {"s01": 0.3}, seed=5), model="fake", seed=5, n_resamples=TEST_RESAMPLES)
    chosen = select_candidates(result, limit=4)
    assert len({s.id for s in chosen}) == len(chosen)


# --- stage 2 ----------------------------------------------------------------


def test_confirmation_shares_one_control_arm():
    """The full-config arm is byte-identical for every candidate, so re-running
    it per comparison would buy the same trials at three times the price."""
    sections = _sections(6)
    agent = SimulatedAgent(sections, {}, seed=1)
    confirm(sections, sections[:3], agent, trials=5)
    n_variants, trials, stage = agent.calls[0]
    assert (n_variants, trials, stage) == (4, 5, "confirmatory")  # 1 baseline + 3 treatments


def test_confirmation_detects_a_section_that_matters():
    sections = _sections(6)
    agent = SimulatedAgent(sections, {"s03": 0.45}, seed=7)
    confirmations = confirm(sections, [sections[2]], agent, trials=5, seed=1)
    assert confirmations[0].pass_rate_verdict is Verdict.REGRESSION
    assert not confirmations[0].safe_to_delete


def test_confirmation_clears_a_section_that_does_nothing():
    """EQUIVALENT requires a margin the design can actually resolve.

    At 12 tasks x 8 trials the MDE is ~21pp, so a +/-20pp margin sits exactly
    at the resolution limit and INCONCLUSIVE would be the honest answer there.
    The margin here is deliberately above it, so the test exercises the
    equivalence path rather than the underpowered one.
    """
    sections = _sections(6)
    agent = SimulatedAgent(sections, {}, seed=8)
    confirmations = confirm(
        sections, [sections[4]], agent, trials=8, margin_overrides={"pass_rate": 0.30}, seed=1
    )
    assert confirmations[0].pass_rate_verdict is Verdict.EQUIVALENT
    assert confirmations[0].safe_to_delete


def test_inconclusive_never_licenses_a_deletion():
    """"We didn't look hard enough" is not permission to delete."""
    sections = _sections(6)
    agent = SimulatedAgent(sections, {}, seed=9)
    confirmations = confirm(
        sections, [sections[1]], agent, trials=5, margin_overrides={"pass_rate": 0.001}, seed=1
    )
    assert confirmations[0].pass_rate_verdict is Verdict.INCONCLUSIVE
    assert not confirmations[0].safe_to_delete


def test_no_candidates_means_no_confirmation_runs():
    sections = _sections(4)
    agent = SimulatedAgent(sections, {}, seed=1)
    assert confirm(sections, [], agent, trials=5) == []
    assert agent.calls == []


# --- report -----------------------------------------------------------------


def _result(effects=None, seed=11, pinned=None, limit=3) -> AblationResult:
    sections = _sections(6, pinned=pinned)
    agent = SimulatedAgent(sections, effects or {}, seed=seed)
    screening = screen(sections, agent, model="haiku-4.5", seed=seed, n_resamples=TEST_RESAMPLES)
    candidates = select_candidates(screening, limit=limit)
    confirmations = confirm(
        sections, candidates, agent, trials=5, margin_overrides={"pass_rate": 0.20}, seed=seed
    )
    return AblationResult(
        sections=tuple(sections),
        screening=screening,
        confirmations=tuple(confirmations),
        config_path="CLAUDE.md",
        total_tokens=sum(s.tokens for s in sections),
    )


def _render(result: AblationResult) -> str:
    buffer = io.StringIO()
    render_ablation(Console(file=buffer, width=100, no_color=True), result, version="0.3.0")
    return buffer.getvalue()


def test_report_labels_screening_as_screening():
    output = _render(_result())
    assert "not confirmatory" in output
    assert "ranking, not the magnitudes" in output


def test_report_shows_token_cost_on_every_screening_row():
    result = _result()
    output = _render(result)
    assert "tokens" in output
    for row in result.screening.rows:
        assert row.section.id in output


def test_report_names_unconfirmed_sections_rather_than_dropping_them():
    """Omitting them would be a lie by selection."""
    result = _result(limit=2)
    output = _render(result)
    assert "Remaining" in output
    for row in result.unscreened:
        assert row.section.id in output


def test_report_states_the_empty_config_floor():
    assert "Empty-config floor" in _render(_result())


def test_report_recommends_only_deletions():
    output = _render(_result())
    assert "Delete" in output or "No section was confirmed safe to delete" in output
    assert "add a section" not in output.lower()


def test_report_prints_the_apply_command_when_there_is_something_to_delete():
    result = _result()
    output = _render(result)
    if result.deletions:
        assert "ablate apply --drop" in output


def test_report_marks_pinned_sections():
    assert "pinned" in _render(_result(pinned={0}))


def test_report_names_the_method():
    output = _render(_result())
    assert "randomized subset screening" in output
    assert "Benjamini-Hochberg" in output


def test_report_survives_a_run_with_no_confirmations():
    result = _result(limit=0)
    output = _render(result)
    assert "SCREENING" in output
    assert "No section was confirmed safe to delete" in output
