from typing import Annotated, Any

import torch
from pydantic import BaseModel, Field, PlainSerializer

from deception_detection.rollouts.api import grade_single_dialogue_for_deception
from deception_detection.types import GradingPromptGetter, Judgement, Label, Message

device = "cuda" if torch.cuda.is_available() else "cpu"


class RolloutDialogue(BaseModel):
    """
    Represents several rollouts for a single input dialogue, together with model grading of outputs.

    This class encapsulates the components of a dialogue interaction, including
    input and output messages, honesty expectations, and model-based evaluations.

    Attributes:
        input_messages (list[Message]): The input messages to the model.
        expect_deception (bool): Indicates whether deception is expected in the responses.
        input_str (str | None): An optional string representation of the input that is
            exactly what is passed to the model.
        output_messages (list[Message]): The generated response messages.
        output_strs (list[str]): String representations of the output messages.
        grader_reasoning (list[str]): The full output of the model grader.
        graded_deception (list[Judgement]): Honesty judgements made by the model grader.
        deception_percentage (float | None): The overall deception score over all rollouts for
            this input dialogue.
    """

    input_messages: list[Message]
    expected_label: Annotated[Label, PlainSerializer(lambda v: v.value)]
    input_str: str | None = None
    output_messages: list[Message] = Field(default_factory=list)
    output_strs: list[str] = Field(default_factory=list)
    grader_reasoning: list[str] = Field(default_factory=list)
    graded_deception: list[Judgement] = Field(default_factory=list)
    deception_percentage: float | None = None
    metadata: dict[str, Any] | None = None

    def grade_all(
        self,
        grading_prompt_getter: GradingPromptGetter,
        verbose: bool,
        mock: bool,
    ) -> None:
        for output_message in self.output_messages:
            grade, response = grade_single_dialogue_for_deception(
                grading_prompt_getter,
                input_messages=self.input_messages,
                output_message=output_message,
                metadata=self.metadata,
                verbose=verbose,
                mock=mock,
            )
            self.graded_deception.append(grade)
            self.grader_reasoning.append(response)
        self.deception_percentage = sum(j == "deceptive" for j in self.graded_deception) / len(
            self.graded_deception
        )
