from dataclasses import dataclass
from enum import Enum
from typing import Any, Literal, Protocol

Judgement = Literal["honest", "deceptive", "ambiguous", "skip"] | str | int

Role = Literal["system", "user", "assistant"]


@dataclass
class Message:
    """
    We allow for messages from the same role to be adjacent.
    In this case, such messages are concatenated.
    """

    role: Role
    content: str
    detect: bool


Dialogue = list[Message]


class GradingPromptGetter(Protocol):
    def __call__(
        self, input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None
    ) -> dict[str, str]: ...

    """
    Get a prompt for grading a given dialogue and response.
    Generally defined by `DialogueDataset.get_grading_prompt`.
    """


class Label(Enum):
    HONEST = "honest"
    DECEPTIVE = "deceptive"
    AMBIGUOUS = "ambiguous"
    SKIP = "skip"


# type alias for the output of `_get_dialogues()`
# dialogues, labels, metadata
# todo: convert to named tuple?
DialogueDataType = tuple[list[Dialogue], list[Label], dict[str, list[Any]] | None]
