from __future__ import annotations

from pathlib import Path

from dovla_cil.data.sharding import iter_cil_records
from dovla_cil.generation.pipeline import generate_builtin_toy_dataset, generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.utils.io import iter_jsonl, read_json


def test_toy_generation_pipeline_writes_shards_and_group_index(tmp_path: Path) -> None:
    summary = generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=tmp_path,
        num_states_per_task=1,
        k=4,
        seed=11,
        shard_size=4,
        inline_observations=True,
    )

    manifest = read_json(tmp_path / "manifest.json")
    group_index = list(iter_jsonl(tmp_path / "group_index.jsonl"))
    records = []
    for shard in manifest["shards"]:
        records.extend(iter_cil_records(tmp_path / str(shard["path"])))

    assert summary.num_groups == 2
    assert summary.num_records == 8
    assert manifest["record_count"] == 8
    assert manifest["group_count"] == 2
    assert manifest["group_index_path"] == "group_index.jsonl"
    assert len(group_index) == 2
    assert all(entry["num_records"] == 4 for entry in group_index)
    assert all(record.rank_within_group is not None for record in records)
    assert all(record.regret is not None for record in records)
    assert any(record.reward.terminal_success for record in records)
    assert "expert" in summary.candidate_type_distribution


def test_builtin_group_shortcut_does_not_cap_at_library_size(tmp_path: Path) -> None:
    summary = generate_builtin_toy_dataset(
        out_dir=tmp_path,
        groups=12,
        k=2,
        seed=17,
        shard_size=64,
    )

    assert summary.num_groups == 12
    assert summary.num_records == 24