from __future__ import annotations

from pathlib import Path

from dovla_cil.data.datasets import CILDataset, write_cil_collection
from dovla_cil.data.group_sampler import GroupAwareBatchSampler
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.training.collate import collate_cil_records


def _make_toy_dataset(tmp_path: Path) -> CILDataset:
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:3],
        out_dir=tmp_path,
        num_states_per_task=1,
        k=4,
        seed=13,
        shard_size=8,
        inline_observations=True,
    )
    return CILDataset(tmp_path)


def test_cil_dataset_loads_generated_toy_cil(tmp_path: Path) -> None:
    dataset = _make_toy_dataset(tmp_path)
    assert len(dataset) == 12
    first = dataset[0]
    assert dataset.get_group(first.group_id)[0] == first
    assert len(list(dataset.iter_groups())) == 3


def test_full_group_sampler_keeps_group_ids_together(tmp_path: Path) -> None:
    dataset = _make_toy_dataset(tmp_path)
    sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1)
    batch_indices = next(iter(sampler))
    records = [dataset[index] for index in batch_indices]
    assert len({record.group_id for record in records}) == 1
    assert len(records) == 4


def test_pair_sampler_returns_same_group_reward_ordered_pairs(tmp_path: Path) -> None:
    dataset = _make_toy_dataset(tmp_path)
    sampler = GroupAwareBatchSampler(
        dataset,
        mode="pairs",
        batch_groups=3,
        pair_count_per_group=2,
        shuffle=False,
        seed=3,
    )
    batch_indices = next(iter(sampler))
    records = [dataset[index] for index in batch_indices]
    assert batch_indices.pair_indices
    for better_index, worse_index in batch_indices.pair_indices:
        better = records[better_index]
        worse = records[worse_index]
        assert better.group_id == worse.group_id
        assert better.reward.score > worse.reward.score


def test_collate_returns_tensors_and_metadata(tmp_path: Path) -> None:
    dataset = _make_toy_dataset(tmp_path)
    sampler = GroupAwareBatchSampler(dataset, mode="full_group", batch_groups=1)
    records = [dataset[index] for index in next(iter(sampler))]
    batch = collate_cil_records(records)

    assert batch["observations"].shape[0] == len(records)
    assert batch["action_features"].shape[0] == len(records)
    assert batch["effect_features"].shape[0] == len(records)
    assert batch["rewards"].shape == (len(records),)
    assert len(batch["instructions"]) == len(records)
    assert len(batch["action_chunks"]) == len(records)
    assert len(batch["effects"]) == len(records)
    assert len(batch["group_ids"]) == len(records)
    assert len(batch["candidate_types"]) == len(records)
    assert len(batch["failures"]) == len(records)
    assert "pair_indices" in batch


def test_zero_copy_collection_loads_disjoint_source_datasets(tmp_path: Path) -> None:
    tasks = built_in_toy_tasks()
    sources = [tmp_path / "source-a", tmp_path / "source-b"]
    for source, task, seed in zip(sources, tasks[:2], (31, 47), strict=False):
        generate_cil_dataset(
            backend="toy",
            tasks=[task],
            out_dir=source,
            num_states_per_task=2,
            k=3,
            seed=seed,
            shard_size=8,
            inline_observations=True,
        )
    collection_dir = tmp_path / "collection"
    write_cil_collection(collection_dir, sources, dataset_name="two-task-collection")

    dataset = CILDataset(collection_dir)

    assert len(dataset) == 12
    assert len(dataset.group_ids) == 4
    assert dataset.index.metadata["dataset_name"] == "two-task-collection"
    assert dataset.index.metadata["task_count"] == 2
    assert {record.metadata["source_dataset"] for record in dataset} == {
        str(source.resolve()) for source in sources
    }