import uuid
from typing import Any, Optional

from genderbench.probing.attempt import Attempt
from genderbench.probing.prompt import Prompt


class ProbeItem:
    """`ProbeItem` is a single test item for `Probe`. Its main role is to group
    related `Attempts` into a single evaluated entity when metrics are
    calculated. The sampling process performed during confidence interval
    calculations is also done on a `ProbeItem` level. For this reason, it is
    recommended to group related `Prompts` into a single `ProbeItem` to avoid
    overconfident inverval ranges.

    `ProbeItem` object handles the creation of multiple `Attempts` with
    identical `Prompts` when `num_repetitions` is used in `Probe`.

    Note:

        Apart from repetitions, `ProbeItem` can be used to group very similar
        prompts that we do not want to necessarily treat as separate entities.
        For example, multiple orderings of options for a
        multiple-choice question can be grouped within one `ProbeItem`.

    Args:
        prompts (list[Prompt]): `Prompts` that will be used to create new
            `Attempts`.
        num_repetitions (int): How many `Attempts` will be created for each
            `Prompt`.
        metadata (Optional[dict[str, Any]], optional): Metadata related to `ProbeItem`
            that can be used during metric calculation process.

    Attributes:
        uuid (uuid.UUID): UUID identifier.
        attempts (list[Attempt]): The list of `Attempts` belonging to
            `ProbeItem`.

    """

    def __init__(
        self,
        prompts: list[Prompt],
        num_repetitions: int,
        metadata: Optional[dict[str, Any]] = None,
    ) -> None:
        self.prompts = prompts
        self.num_repetitions = num_repetitions
        self.metadata = metadata
        self.uuid: uuid.UUID = uuid.uuid4()

        self.attempts: list[Attempt] = [
            Attempt(prompt, repetition_id)
            for prompt in self.prompts
            for repetition_id in range(self.num_repetitions)
        ]

    def to_json_dict(self) -> dict:
        """Prepare a JSON-serializable dictionary representation. Used for
        logging.

        Returns:
            dict: JSON-serializable dictionary.
        """
        parameters = ["uuid", "num_repetitions", "metadata"]
        d = {parameter: getattr(self, parameter) for parameter in parameters}
        d["prompts"] = [prompt.to_json_dict() for prompt in self.prompts]
        d["attempts"] = [attempt.to_json_dict() for attempt in self.attempts]
        return d

    @classmethod
    def from_json_dict(cls, json_dict: dict) -> "ProbeItem":
        """Create a new `ProbeItem` object from a JSON-serializable dictionary
        representation.

        Args:
            json_dict (dict): JSON-serializable dictionary. Generated by
                ``to_json_dict``.

        Returns:
            Attempt: Restored `ProbeItem` object.
        """
        prompts = [Prompt.from_json_dict(prompt) for prompt in json_dict["prompts"]]
        prompt_uuids = {str(prompt.uuid): prompt for prompt in prompts}
        item = ProbeItem(
            prompts=prompts,
            num_repetitions=json_dict["num_repetitions"],
            metadata=json_dict["metadata"],
        )
        item.uuid = uuid.UUID(json_dict["uuid"])
        item.attempts = [
            Attempt.from_json_dict(attempt, prompt_uuids[attempt["prompt"]["uuid"]])
            for attempt in json_dict["attempts"]
        ]
        return item
