import pytest

from deception_detection import data
from deception_detection.data import DialogueDataset
from deception_detection.repository import ALL_DATASETS_DICT


@pytest.mark.parametrize("dataset_id", ALL_DATASETS_DICT.keys())
def test_load_all_datasets(dataset_id: str):
    dataset_factory = ALL_DATASETS_DICT[dataset_id]
    dataset = dataset_factory()
    assert isinstance(dataset, DialogueDataset)
    assert len(dataset.dialogues) > 0
    assert len(dataset.labels) == len(dataset.dialogues)
    if dataset.metadata is not None:
        assert isinstance(dataset.metadata, dict)
        for k, v in dataset.metadata.items():
            assert isinstance(v, list)
            assert len(v) == len(dataset.dialogues), f"Metadata {k} has incorrect length"


@pytest.mark.parametrize("dataset_id", ALL_DATASETS_DICT.keys())
def test_dataset_slicing(dataset_id: str):
    dataset_factory = ALL_DATASETS_DICT[dataset_id]
    dataset = dataset_factory()

    # Test slicing from the beginning
    slice_len = 5
    sliced_dataset = dataset[:slice_len]
    assert isinstance(sliced_dataset, type(dataset))
    assert len(sliced_dataset) == min(slice_len, len(dataset))

    # Test slicing from the end
    sliced_dataset = dataset[-slice_len:]
    assert isinstance(sliced_dataset, type(dataset))
    assert len(sliced_dataset) == min(slice_len, len(dataset))

    # Test slicing with a step
    sliced_dataset = dataset[::2]
    assert isinstance(sliced_dataset, type(dataset))
    assert len(sliced_dataset) == (len(dataset) + 1) // 2

    # Test slicing with start and end
    mid = len(dataset) // 2
    sliced_dataset = dataset[mid - 2 : mid + 2]
    assert isinstance(sliced_dataset, type(dataset))
    assert len(sliced_dataset) == min(4, len(dataset))

    # Test that sliced dataset preserves attributes
    assert sliced_dataset.base_name == dataset.base_name
    assert sliced_dataset.variant == dataset.variant
    if hasattr(dataset, "metadata"):
        assert hasattr(sliced_dataset, "metadata")
    if hasattr(dataset, "grading_function"):
        assert hasattr(sliced_dataset, "grading_function")


@pytest.mark.parametrize("dataset_class", ALL_DATASETS_DICT.values())
def test_add_datasets(dataset_class: type[DialogueDataset]):
    dataset = dataset_class()

    # Test adding dataset to itself
    add_to_self = dataset + dataset
    assert isinstance(add_to_self, DialogueDataset)
    assert len(add_to_self) == len(dataset) * 2

    # Test adding dataset to another dataset
    other_dataset = next((d() for d in ALL_DATASETS_DICT.values() if d != dataset_class), None)
    if other_dataset:
        add_to_other = dataset + other_dataset
        assert isinstance(add_to_other, DialogueDataset)
        assert len(add_to_other) == len(dataset) + len(other_dataset)
        assert add_to_other.base_name == f"{dataset.base_name}+{other_dataset.base_name}"
        assert add_to_other.variant == f"({dataset.variant}+{other_dataset.variant})"


# commented out as these datasets arenot supported for now


@pytest.mark.parametrize(
    "dataset_class,metadata_keys",
    [
        # (InsiderTradingPrewritten, ["difficulties", "prompt_variation"]),
        # (PoliticalSycophancyDataset, ["syco_score"]),
        # (PoetryFlattery, ["reward", "passed"]),
        # (InsubordinateRubricModification, ["reward", "passed"]),
        # (NudgedRubricModification, ["reward", "passed"]),
        # (RewardTampering, ["reward", "passed"]),
        (data.WerewolfDataset, {"role", "night_info", "discussion"}),
        (data.RolePlayingDataset, {"scenario", "question", "answer_prefix"}),
    ],
)
def test_metadata(dataset_class: type[DialogueDataset], metadata_keys: set[str]):
    dataset = dataset_class(variant=dataset_class.default_variant)
    assert dataset.metadata is not None
    assert set(dataset.metadata.keys()) == metadata_keys
    for key in metadata_keys:
        assert len(dataset.metadata[key]) == len(dataset.dialogues)

    sliced_dataset = dataset[:10]
    assert sliced_dataset.metadata is not None
    assert set(sliced_dataset.metadata.keys()) == metadata_keys
    for key in metadata_keys:
        assert len(sliced_dataset.metadata[key]) == len(sliced_dataset.dialogues)
