from pathlib import Path
from typing import Any, List

from huggingface_hub import snapshot_download
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.model import (
    ChatMessage,
    ChatMessageSystem,
    ChatMessageUser,
    ContentImage,
    ContentText,
)
from platformdirs import user_cache_dir

DATA_BASE_DIR = Path(user_cache_dir("inspect_evals")) / "cyse3_visual_prompt_injection"

HF_PATH = "facebook/cyberseceval3-visual-prompt-injection"


def record_to_sample(record: dict[str, Any]) -> Sample:
    messages: list[ChatMessage] = [
        ChatMessageSystem(content=record["system_prompt"]),
        ChatMessageUser(
            content=[
                ContentText(text=record["user_input_text"]),
                ContentImage(image=str(DATA_BASE_DIR / f"images/{record['id']}.png")),
            ]
        ),
    ]

    return Sample(
        input=messages,
        target=record["judge_question"],
        metadata={
            "id": record["id"],
            "system_prompt": record["system_prompt"],
            "user_input_text": record["user_input_text"],
            "judge_question": record["judge_question"],
            "injection_technique": record["injection_technique"],
        },
    )


def read_dataset(
    challenges: str | List[str] | None = None,
    shuffle: bool = False,
) -> Dataset:
    """Reads a dataset from a test_cases.json file."""
    snapshot_download(HF_PATH, repo_type="dataset", local_dir=DATA_BASE_DIR)

    dataset = hf_dataset(
        HF_PATH,
        split="test",
        sample_fields=record_to_sample,
        shuffle=shuffle,
    )

    if challenges:
        challenges = challenges if isinstance(challenges, list) else [challenges]
        challenge_ids = [
            sample.metadata["injection_technique"]
            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_technique"] in challenges
        )

    return dataset
