from __future__ import annotations

import json
from pathlib import Path

import pytest

from eval.score import score_run

pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]


@pytest.mark.parametrize(
    ("sample_indices", "complete"),
    [
        pytest.param([[0, 1], [0, 1]], True, id="complete"),
        pytest.param([[1, 0], [1, 0]], True, id="out-of-order"),
        pytest.param([[0], [0, 1, 2]], False, id="extra-sample-hides-missing-sample"),
        pytest.param([[1, 2], [0, 1]], False, id="shifted-sample-indices"),
        pytest.param([[-1, 0], [0, 1]], False, id="negative-sample-index"),
        pytest.param([[0], [0, 1]], False, id="missing-sample"),
        pytest.param([[], [0, 1, 2, 3]], False, id="missing-task"),
        pytest.param([[0, 1, 2], [0, 1]], False, id="extra-sample"),
    ],
)
def test_score_run_sample_coverage(
    tmp_path: Path, capsys: pytest.CaptureFixture[str], sample_indices: list[list[int]], complete: bool
) -> None:
    data_file = tmp_path / "data.jsonl"
    data_file.write_text(
        "".join(json.dumps({"task_id": str(i), "question": "What is 1 + 1?", "answer": "2"}) + "\n" for i in range(2)),
        encoding="utf-8",
    )
    (tmp_path / "meta.json").write_text(json.dumps({"bench": "aime", "data_file": str(data_file), "n_sample": 2}), encoding="utf-8")
    for task_id, indices in enumerate(sample_indices):
        (tmp_path / f"raw_rank{task_id:03d}.jsonl").write_text(
            "".join(json.dumps({"task_id": str(task_id), "sample_idx": index, "raw_text": r"\boxed{2}"}) + "\n" for index in indices),
            encoding="utf-8",
        )

    assert score_run(str(tmp_path), judge="rule", verbose=False) == 0

    summary = json.loads((tmp_path / "summary.json").read_text(encoding="utf-8"))
    assert summary["coverage_complete"] is complete
    assert summary["expected_rollouts"] == 4
    assert summary["num_rollouts"] == sum(map(len, sample_indices))
    assert ("[warn] incomplete coverage" in capsys.readouterr().out) is (not complete)