from typing import Any, List

from inspect_ai.dataset import Dataset, Sample, json_dataset


def record_to_sample(record: dict[str, Any]) -> Sample:
    return Sample(
        input=record["mutated_prompt"],
        metadata={
            "attack_type": record["attack_type"],
        },
    )


def read_dataset(challenges: str | List[str] | None = None) -> Dataset:
    """Reads a dataset from a challenges.json file."""
    dataset = json_dataset(
        json_file="https://raw.githubusercontent.com/meta-llama/PurpleLlama/70e7a376f2a310fb89f7ad207685baa69f78107f/CybersecurityBenchmarks/datasets/interpreter/interpreter.json",
        sample_fields=record_to_sample,
    )

    # filter if requested
    if challenges:
        from itertools import chain

        # get challenges and validate them
        challenges = challenges if isinstance(challenges, list) else [challenges]
        valid_challenges = [
            sample.metadata["attack_type"]
            for sample in dataset
            if sample.metadata is not None
        ]
        valid_challenges = list(set(chain.from_iterable(valid_challenges)))
        invalid_challenges = [
            challenge for challenge in challenges if challenge not in valid_challenges
        ]
        if invalid_challenges:
            challenge_names = ", ".join(invalid_challenges)
            raise ValueError(f"Invalid challenge names: {challenge_names}")

        # apply filter
        dataset = dataset.filter(
            lambda sample: sample.metadata is not None
            and (set(sample.metadata["attack_type"]) & set(challenges)) != set()
        )

    return dataset
