# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import logging
from typing import cast

from pyrit.datasets.seed_datasets.remote.remote_dataset_loader import (
    _RemoteDatasetLoader,
)
from pyrit.models import Modality, SeedDataset, SeedPrompt, SeedUnion
from pyrit.models.harm_category import HarmCategory

logger = logging.getLogger(__name__)


class _ORBenchBaseDataset(_RemoteDatasetLoader):
    """
    Base loader for OR-Bench datasets from HuggingFace.

    Subclasses must set CONFIG, provide a dataset_name property, and a description.

    References:
        - https://huggingface.co/datasets/bench-llm/OR-Bench
        - [@cui2024orbench]
    License: CC BY 4.0

    Warning: This dataset contains prompts designed to test over-refusal behavior in LLMs,
    including potentially harmful and toxic content.
    """

    HF_DATASET_NAME: str = "bench-llm/OR-Bench"
    CONFIG: str
    DESCRIPTION: str
    # or-bench-80k and or-bench-hard-1k are BENIGN over-refusal prompts: their `category`
    # names the harm domain the *safe* prompt superficially resembles, not an actual harm,
    # so harm_categories is left empty. Only or-bench-toxic contains genuinely harmful
    # prompts, so only that subset maps `category` to the canonical taxonomy.
    MAPS_HARM_CATEGORIES: bool = False
    HARM_ALIAS_OVERRIDES: dict[str, list[HarmCategory]] = cast("dict[str, list[HarmCategory]]", {})

    should_register = False  # abstract base — subclasses register themselves

    # Metadata shared across all OR-Bench subclasses; subclasses override `size`.
    modalities: tuple[Modality, ...] = (Modality.TEXT,)
    tags: frozenset[str] = frozenset({"default", "safety", "refusal"})

    async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset:
        """
        Fetch OR-Bench dataset from HuggingFace and return as SeedDataset.

        Args:
            cache: Whether to cache the fetched dataset. Defaults to True.

        Returns:
            SeedDataset: A SeedDataset containing the OR-Bench prompts.
        """
        logger.info(f"Loading OR-Bench dataset from {self.HF_DATASET_NAME} (config={self.CONFIG})")

        data = await self._fetch_from_huggingface_async(
            dataset_name=self.HF_DATASET_NAME,
            config=self.CONFIG,
            split="train",
            cache=cache,
        )

        authors = [
            "Justin Cui",
            "Wei-Lin Chiang",
            "Ion Stoica",
            "Cho-Jui Hsieh",
        ]
        source_url = f"https://huggingface.co/datasets/{self.HF_DATASET_NAME}"
        groups = ["UCLA", "UC Berkeley"]

        seed_prompts: list[SeedUnion] = [
            SeedPrompt(
                value=item["prompt"],
                data_type="text",
                dataset_name=self.dataset_name,
                harm_categories=(
                    self._standardize_harm_categories(item.get("category"), alias_overrides=self.HARM_ALIAS_OVERRIDES)
                    if self.MAPS_HARM_CATEGORIES
                    else []
                ),
                description=self.DESCRIPTION,
                source=source_url,
                authors=authors,
                groups=groups,
                metadata={"category": item.get("category", "")},
            )
            for item in data
        ]

        logger.info(f"Successfully loaded {len(seed_prompts)} prompts from OR-Bench dataset")

        return SeedDataset(seeds=seed_prompts, dataset_name=self.dataset_name)


class _ORBench80KDataset(_ORBenchBaseDataset):
    """
    Loader for the OR-Bench 80K dataset.

    Contains ~80k over-refusal prompts categorized into 10 common rejection categories.
    This is the main comprehensive benchmark for evaluating LLM over-refusal behavior.
    """

    CONFIG: str = "or-bench-80k"
    DESCRIPTION: str = (
        "OR-Bench 80K contains ~80k over-refusal prompts categorized into 10 rejection "
        "categories. This is the main comprehensive benchmark for evaluating LLM over-refusal."
    )
    size: str = "huge"  # 80359 over-refusal prompts
    should_register = True

    @property
    def dataset_name(self) -> str:
        """The dataset name."""
        return "or_bench_80k"


class _ORBenchHardDataset(_ORBenchBaseDataset):
    """
    Loader for the OR-Bench Hard-1K dataset.

    Contains ~1k challenging safe prompts that commonly trigger over-refusal in LLMs.
    These are prompts that models should be able to answer without refusing.
    """

    CONFIG: str = "or-bench-hard-1k"
    DESCRIPTION: str = (
        "OR-Bench Hard-1K contains ~1k challenging safe prompts that commonly trigger "
        "over-refusal in language models. These prompts should be answerable without refusing."
    )
    size: str = "large"  # 1319 challenging safe prompts
    should_register = True

    @property
    def dataset_name(self) -> str:
        """The dataset name."""
        return "or_bench_hard"


class _ORBenchToxicDataset(_ORBenchBaseDataset):
    """
    Loader for the OR-Bench Toxic dataset.

    Contains toxic prompts that language models should correctly refuse.
    Used as a contrast set to evaluate whether models can distinguish
    genuinely harmful prompts from safe ones.
    """

    CONFIG: str = "or-bench-toxic"
    # Unlike the benign 80k/hard subsets, or-bench-toxic prompts are genuinely harmful,
    # so their `category` is a real harm domain and is standardized to the taxonomy.
    MAPS_HARM_CATEGORIES: bool = True
    HARM_ALIAS_OVERRIDES: dict[str, list[HarmCategory]] = {
        "hate": [HarmCategory.HATE_SPEECH, HarmCategory.REPRESENTATIONAL],
        "privacy": [HarmCategory.PPI],
        "harmful": [HarmCategory.OTHER],
        "unethical": [HarmCategory.OTHER],
    }
    DESCRIPTION: str = (
        "OR-Bench Toxic contains toxic prompts that language models should correctly refuse. "
        "Used as a contrast set to evaluate refusal calibration."
    )
    size: str = "large"  # 655 toxic prompts for refusal calibration
    should_register = True

    @property
    def dataset_name(self) -> str:
        """The dataset name."""
        return "or_bench_toxic"