"""Tests for the rollout assembly seam (fake verifier; no model, no engine)."""

import enum

import pytest
import torch

from align.rl.advantages import GroupFilterStats
from align.rl.assembly import assemble_rollout
from align.rl.rewards import CompositeReward


class _Outcome(enum.Enum):
    CORRECT = 1
    WRONG = 2


def _verify(golds, preds):
    return [
        _Outcome.CORRECT if (p is not None and p.strip() == g) else _Outcome.WRONG
        for g, p in zip(golds, preds)
    ]


def _text(answer: str) -> str:
    return f"Step 1: reason.\n\n# Answer\n\n{answer}"


def _mk(n_prompts=2, g=2, answers=None):
    """Two prompts x G completions with controllable answers."""
    prompt_ids = [[1, 2, 3], [4, 5]]
    completions = [[[10, 11], [12, 13, 14]], [[20], [21, 22]]]
    logprobs = [[[-0.1, -0.2], [-0.3, -0.4, -0.5]], [[-0.6], [-0.7, -0.8]]]
    answers = answers or [["7", "8"], ["9", "9"]]
    texts = [[_text(a) for a in row] for a in [None] for row in answers]
    return prompt_ids, completions, logprobs, texts


def test_assembles_shapes_masks_and_alignment():
    prompt_ids, completions, logprobs, texts = _mk(answers=[["7", "8"], ["9", "1"]])
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is not None
    B, T = out.input_ids.shape
    assert B == 4 and T == 6  # longest row: prompt 3 + completion 3
    # row 0: prompt [1,2,3] + completion [10,11]
    assert out.input_ids[0].tolist() == [1, 2, 3, 10, 11, 0]
    assert out.response_mask[0].tolist() == [0, 0, 0, 1, 1, 0]
    # label-aligned engine logprobs: completion token t -> position len(p)+t-1
    assert out.vllm_logprobs[0, 2].item() == pytest.approx(-0.1)
    assert out.vllm_logprobs[0, 3].item() == pytest.approx(-0.2)
    assert out.vllm_logprobs[0, 1].item() == 0.0  # prompt region zero-filled
    # group ids: two kept groups
    assert out.group_ids.tolist() == [0, 0, 1, 1]
    # advantages: group means subtracted -> each group sums to ~0
    a = out.advantages
    assert abs(float(a[0] + a[1])) < 1e-6 and abs(float(a[2] + a[3])) < 1e-6
    # rewards itemised
    assert set(out.rewards) >= {"reward/accuracy", "reward/format", "reward/total"}
    assert out.rewards["reward/accuracy"].tolist() == [1.0, 0.0, 1.0, 0.0]


def test_dead_groups_dropped_whole_with_telemetry():
    # group 0: both wrong (all_zero-ish variance 0 at total level? totals equal)
    # -> dead; group 1: mixed -> kept
    prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["9", "2"]])
    stats = GroupFilterStats(max_possible_score=1.1)  # acc 1.0 + fmt 0.1
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
        filter_stats=stats,
    )
    assert out is not None
    assert out.kept_groups == 1 and out.dropped_groups == 1
    assert out.input_ids.shape[0] == 2  # only group 1's two rows
    # both-wrong completions still earn the format term (0.1 each), so the dead
    # group's identical totals are NONZERO -> nonzero_std0, not all_zero. With a
    # welded soft-format reward, all_zero is reserved for unformatted garbage.
    assert stats.nonzero_std0 == 1 and stats.kept == 1


def test_all_dead_returns_none_not_error():
    prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["2", "2"]])
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is None


def test_ragged_logprobs_refused():
    prompt_ids, completions, logprobs, texts = _mk()
    logprobs[0][0] = [-0.1]  # 1 logprob for 2 completion tokens
    with pytest.raises(ValueError, match="ragged|logprobs"):
        assemble_rollout(
            prompt_ids,
            completions,
            logprobs,
            texts,
            golds=["7", "9"],
            reward=CompositeReward(format_weight=0.1),
            verify_batch=_verify,
            pad_token_id=0,
        )


def test_unanswerable_text_scores_zero_accuracy():
    prompt_ids, completions, logprobs, _ = _mk()
    texts = [["no answer marker here", _text("8")], [_text("9"), _text("9")]]
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["8", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is not None
    # group 0 kept (one right, one no-answer); the no-answer row scores acc 0
    row0_acc = out.rewards["reward/accuracy"][0].item()
    assert row0_acc == 0.0