from typing import Any, List

from inspect_ai.dataset import Dataset, Sample, json_dataset

DATASET_URL = "https://raw.githubusercontent.com/meta-llama/PurpleLlama/70e7a376f2a310fb89f7ad207685baa69f78107f/CybersecurityBenchmarks/datasets/prompt_injection/prompt_injection.json"


def record_to_sample(record: dict[str, Any]) -> Sample:
    return Sample(
        input=record["user_input"],
        target=record["judge_question"],
        metadata={
            "test_case_prompt": record["test_case_prompt"],
            "judge_question": record["judge_question"],
            "injection_variant": record["injection_variant"],
        },
    )


def read_dataset(
    challenges: str | List[str] | None = None,
    shuffle: bool = False,
) -> Dataset:
    """Reads a dataset from a challenges.json file."""
    dataset = json_dataset(
        json_file=DATASET_URL,
        sample_fields=record_to_sample,
        shuffle=shuffle,
    )

    if challenges:
        challenges = challenges if isinstance(challenges, list) else [challenges]
        challenge_ids = [
            sample.metadata["injection_variant"]
            for sample in dataset
            if sample.metadata is not None
        ]
        invalid_challenges = [
            challenge for challenge in challenges if challenge not in challenge_ids
        ]
        if invalid_challenges:
            challenge_names = ", ".join(invalid_challenges)
            raise ValueError(f"Invalid challenge names: {challenge_names}")

        dataset = dataset.filter(
            lambda sample: sample.metadata is not None
            and sample.metadata["injection_variant"] in challenges
        )

    return dataset
