"""Attack family registry."""

from __future__ import annotations

from typing import TYPE_CHECKING

from rewardprobe.families.basic_sanity import BasicSanityFamily
from rewardprobe.families.composition import CompositionFamily
from rewardprobe.families.parser_bugs import ParserBugsFamily
from rewardprobe.families.schema import SchemaFamily
from rewardprobe.families.stateful import StatefulFamily
from rewardprobe.families.trivial_bypass import TrivialBypassFamily

if TYPE_CHECKING:
    from rewardprobe.families.base import AttackFamily

FAMILIES: dict[str, type[AttackFamily]] = {
    "basic_sanity": BasicSanityFamily,
    "trivial_bypass": TrivialBypassFamily,
    "parser_bugs": ParserBugsFamily,
    "composition": CompositionFamily,
    "schema": SchemaFamily,
    "stateful": StatefulFamily,
}

FAMILY_ORDER: list[str] = [
    "basic_sanity",
    "trivial_bypass",
    "parser_bugs",
    "composition",
    "schema",
    "stateful",
]

__all__ = [
    "FAMILIES",
    "FAMILY_ORDER",
    "BasicSanityFamily",
    "TrivialBypassFamily",
    "ParserBugsFamily",
    "CompositionFamily",
    "SchemaFamily",
    "StatefulFamily",
]
