from abc import abstractmethod
from collections import defaultdict
from typing import Any, Self

import numpy as np

from deception_detection.log import logger
from deception_detection.models import get_llama3_tokenizer
from deception_detection.types import Dialogue, DialogueDataType, Judgement, Label, Message
from deception_detection.utils import preprocess_dialogue


class DialogueDataset:
    """
    Text dataset containing a list of dialogues. Parent class for PrewrittenDialogueDataset
    and RolloutDialogueDataset.

    `base_name` is the name of the dataset. It is a class attribute and should be set on the
    subclass.

    `all_variants` is a list of all possible variants of the dataset. It is a class attribute and
    should be set on the subclass.

    `padding` specifies how many previous tokens back to detect and how many subsequent tokens to
    detect. This defaults to 4 for gemma, and 2 for mistral. This is because we usually
    detect on whole assistant messages, and would like to also detect on the the various seperator
    tokens.

    `dialogues` is a list of dialogues.

    `is_deceptive` is a list of booleans, where each boolean indicates whether the corresponding
    dialogue is deceptive.

    `metadata` is an optional attribute that can be used to store any additional information about
    the dataset in the form of a dictionary with list of values.

    `grading_function` is an optional attribute that takes in an input dialogue and output message
    and returns a dictionary of the following form, which is passed to the Anthropic API for model
    grading of rollouts from the dataset.

    {
        "system": str
        "user": str
    }


    See `ai_liar.py` for an example of how to implement this class.
    """

    # -- Class attributes --
    base_name: str = "base_name_not_set"
    default_variant: str = "plain"

    # defaults to only "plain" - subclasses can override
    all_variants: list[str] = ["plain"]

    # defaults to prewritten. if loading from rollouts, should overwrite.
    model: str = "prewritten"

    # defaults to None. subclasses can override.
    trim_reasoning_str: str | None = None

    # we define the number of tokens in header and end of message boilerplate as padding, so it can
    # be included in detect masks by default
    # TODO: this is kind of jank.
    # ideally this non zero default would just live in the tokenizer class.
    padding: dict[str, dict[str, int]] = {
        "gemma": {"left": 4, "right": 1},
        "mistral": {"left": 2, "right": 1},
        "llama": {"left": 4, "right": 1},
        "qwen": {"left": 4, "right": 1},
    }

    # -- Instance attribute types--
    dialogues: list[Dialogue]
    variant: str
    labels: list[Label]
    metadata: dict[str, list[Any]] | None

    # ---- Abstract methods (must be implemented by subclasses) ----

    @abstractmethod
    def _get_dialogues(self) -> DialogueDataType:
        """
        Abstract method to be implemented by subclasses, returning the dialogues, is_deceptive,
        and metadata.
        """
        pass

    @staticmethod
    @abstractmethod
    def get_grading_prompt(
        input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None = None
    ) -> dict[str, str]:
        raise NotImplementedError("A grading prompt must be implemented by a subclass")

    # ---- End of abstract methods ----

    def __init__(
        self,
        variant: str | None = None,
        dialogues: list[Dialogue] | None = None,
        labels: list[Label] | None = None,
        metadata: dict[str, list[Any]] | None = None,
        shuffle_upon_init: bool | None = None,
        skip_variant_validation: bool = False,
        trim_reasoning: bool = False,
        model: str | None = None,
        max_llama_token_length: int | None = None,
    ):
        """
        Sometimes implemented by subclasses, but usually not.

        This init code can be called in two ways:
        - Either, we can pass in none of the arguments, in which case the init code will call
        `_get_dialogues()` to get the dialogues, label, and metadata.
        - Or, we can pass in all of the arguments, in which case the init code will not call
        `_get_dialogues()`, and will just set the attributes directly. This is useful for e.g.
        slicing a dataset or otherwise getting some subset of the dialogues, which a user adding
        a new dataset should never have to worry about.
        """

        if variant is None:
            variant = self.default_variant
        if variant not in self.all_variants and not skip_variant_validation:
            raise ValueError(f"Invalid variant: {variant}. Must be one of {self.all_variants}")
        self.variant = variant

        if dialogues is None:
            assert (
                labels is None and metadata is None
            ), "If dialogues is None, labels and metadata must also be None"
            dialogues, labels, metadata = self._get_dialogues()
        assert labels is not None

        self.dialogues = dialogues
        self.labels = labels
        self.metadata = metadata
        if model is not None:
            self.model = model

        self.post_process()
        if trim_reasoning and self.trim_reasoning_str is not None:
            self.split_dialogues_on_reasoning(self.trim_reasoning_str)

            # we should also set the padding to 0 on the left in these cases
            self.padding = {k: {"left": 0, "right": v["right"]} for k, v in self.padding.items()}

        # this is important because we often want to split the dataset and don't
        # want to only end up with deceptive dialogues in some eval set
        # we only do this if the dataset contains a mixture of honest and deceptive dialogues
        # and we don't do it upon slicing
        if shuffle_upon_init is None:
            shuffle_upon_init = True
        if len(set(self.labels)) > 1 and shuffle_upon_init:
            self.randomly_shuffle_dialogues()
        if max_llama_token_length is not None:
            # get only tokenizer
            self.llama_tokenizer = get_llama3_tokenizer()
            indices_to_keep = self.get_short_seq_indices(max_llama_token_length)
            self.dialogues, self.labels, self.metadata = self._slice_by_index(indices_to_keep)

    @classmethod
    def _get_label_from_judgement(cls, judgement: Judgement) -> Label:
        if isinstance(judgement, str):
            match judgement:
                case "honest":
                    return Label.HONEST
                case "deceptive":
                    return Label.DECEPTIVE
                case "unknown":
                    return Label.SKIP
                case "ambiguous" | "refusal":
                    return Label.AMBIGUOUS
                case "skip":
                    return Label.SKIP
                case _:
                    logger.warning(f"Invalid judgement: {judgement}")
                    return Label.SKIP
        else:
            raise ValueError(
                f"Non-string judgement: {judgement}. "
                f"please overload _{cls.__name__}._get_label_from_judgement"
            )

    @classmethod
    def from_rollouts(
        cls,
        base_name: str,
        variant: str,
        model: str,
        rollouts: list[dict[str, Any]],
        trim_reasoning: bool = False,
        shuffle_upon_init: bool = True,
        max_llama_token_length: int | None = None,  # 2048,
    ) -> "Self":
        """
        Load a DialogueDataset from a rollout JSON file.

        `trim_reasoning` is a boolean indicating whether to trim the reasoning from the output
        message. To specify how to trim, specify a property `trim_reasoning_str` in the subclass.
        """
        dialogues: list[Dialogue] = []
        labels: list[Label] = []
        metadata: dict[str, list[Any]] = defaultdict(list)

        for item in rollouts:
            # a rollout dialogue can have multiple completions, so we iterate over them.
            for i in range(len(item["output_strs"])):
                dialogue = [
                    Message(m["role"], m["content"].strip(), m["detect"])
                    for m in item["input_messages"]
                ]
                last_message = dialogue[-1]
                assert last_message.role == "assistant" and last_message.content == ""
                dialogue[-1] = Message("assistant", item["output_strs"][i].rstrip(), True)

                dialogues.append(dialogue)
                labels.append(cls._get_label_from_judgement(item["graded_deception"][i]))
                metadata["model_grader_reasoning"].append(item["grader_reasoning"][i])
                metadata["model_graded_deception"].append(item["graded_deception"][i])
                if "expected_label" in item:
                    metadata["expected_label"].append(item["expected_label"])

                if item["metadata"] is not None:
                    for k, v in item["metadata"].items():
                        metadata[k].append(v)
        if "id" not in metadata:
            metadata["id"] = list(range(len(dialogues)))

        new = cls(
            variant,
            dialogues,
            labels,
            dict(metadata),
            model=model,
            skip_variant_validation=True,
            trim_reasoning=trim_reasoning,
            shuffle_upon_init=shuffle_upon_init,
            max_llama_token_length=max_llama_token_length,
        )
        new.base_name = base_name
        return new

    def randomly_shuffle_dialogues(self) -> None:
        """
        Randomly shuffles the dialogues.
        """
        np.random.seed(42)
        perm = np.random.permutation(len(self))
        self.dialogues = [self.dialogues[i] for i in perm]
        self.labels = [self.labels[i] for i in perm]
        if self.metadata is not None:
            self.metadata = {k: [v[i] for i in perm] for k, v in self.metadata.items()}

    def split_dialogues_on_reasoning(self, trim_reasoning: str) -> None:
        for i, dialogue in enumerate(self.dialogues):
            final_message = dialogue[-1]
            # make sure this is an assistant message
            assert final_message.role == "assistant"
            # make sure this is a message that should be detected
            assert final_message.detect
            # check if the split_on string is in the final message
            content = final_message.content
            if trim_reasoning in content:
                # split the message into two at the split_on string
                split_idx = content.rindex(trim_reasoning)
                dialogue[-1] = Message(role="assistant", content=content[:split_idx], detect=False)
                dialogue.append(Message(role="assistant", content=content[split_idx:], detect=True))
            else:
                self.labels[i] = Label.SKIP
                logger.warning(
                    f"trim_reasoning separator {trim_reasoning} missing for dialogue {i} in dataset"
                    f" {self.base_name}, skipping"
                )

    def post_process(self) -> None:
        """
        Post-process the dataset.
        """
        pass

    def _slice_by_index(self, indices: int | list[int]) -> DialogueDataType:
        if isinstance(indices, int):
            indices = [indices]
        dialogues = [self.dialogues[i] for i in indices]
        labels = [self.labels[i] for i in indices]
        metadata: dict[str, list[Any]] | None = None
        if self.metadata is not None:
            metadata = {k: [v[i] for i in indices] for k, v in self.metadata.items()}

        return dialogues, labels, metadata

    def subset_where_true(self, condition: list[bool]) -> Self:
        """Filter the dataset by a condition on the metadata."""
        assert len(condition) == len(self)
        indices = [i for i, c in enumerate(condition) if c]
        dialogues, labels, metadata = self._slice_by_index(indices)
        new = self.__class__(
            self.variant, dialogues, labels, metadata, skip_variant_validation=True
        )
        new.base_name = self.base_name
        new.padding = self.padding
        new.model = self.model
        return new

    def get_with_label(self, label: Label) -> Self:
        return self.subset_where_true([_label == label for _label in self.labels])

    def get_honest(self) -> Self:
        return self.get_with_label(Label.HONEST)

    def get_deceptive(self) -> Self:
        return self.get_with_label(Label.DECEPTIVE)

    def __len__(self) -> int:
        return len(self.dialogues)

    def _slice(self, idx: int | slice) -> Self:
        """Slices the dataset, preserving metadata."""
        if isinstance(idx, int):
            dialogues, labels, metadata = self._slice_by_index(idx)
        else:  # slice
            indices = list(range(len(self)))[idx]
            dialogues, labels, metadata = self._slice_by_index(indices)

        new = self.__class__(
            self.variant,
            dialogues,
            labels,
            metadata,
            shuffle_upon_init=False,
            skip_variant_validation=True,
        )
        new.base_name = self.base_name
        new.padding = self.padding
        new.model = self.model
        return new

    def __getitem__(self, idx: int | slice):
        """
        Get a partial dataset or pair of tuples.
        """
        return self._slice(idx)

    def __add__(self, other: "DialogueDataset") -> "DialogueDataset":
        """
        Concatenate two DialogueDataset instances, shuffling the result on init.

        If only one of the datasets has metadata, the metadata is not included in the new dataset.
        """
        dialogues: list[Dialogue] = self.dialogues + other.dialogues
        labels: list[Label] = self.labels + other.labels
        if self.metadata is not None and other.metadata is not None:
            assert set(self.metadata.keys()) == set(other.metadata.keys()), "Metadata key mismatch"
            metadata = {k: self.metadata[k] + other.metadata[k] for k in self.metadata}
        else:
            metadata = None
        new_base_name = f"{self.base_name}+{other.base_name}"
        new_variant = f"({self.variant}+{other.variant})"
        new = DialogueDataset(
            new_variant,
            dialogues,
            labels,
            metadata,
            skip_variant_validation=True,
            shuffle_upon_init=True,
        )
        new.base_name = new_base_name
        new.padding = self.padding
        return new

    def get_short_seq_indices(self, max_llama_token_length: int) -> list[int]:
        preprocessed_dialogues = [preprocess_dialogue(d) for d in self.dialogues]
        formatted_dialogues = self.llama_tokenizer.apply_chat_template(
            preprocessed_dialogues, tokenize=False, add_generation_prompt=False
        )
        indices_to_keep = []
        for i, dialogue in enumerate(formatted_dialogues):
            tokenizer_out = self.llama_tokenizer(
                dialogue,  # type: ignore
                padding=True,
                return_tensors="pt",
                add_special_tokens=False,
                truncation=True,
            )
            tokens = tokenizer_out["input_ids"]
            if len(tokens[0]) <= max_llama_token_length:  # type: ignore
                indices_to_keep.append(i)
        return indices_to_keep

    @property
    def name(self) -> str:
        return f"{self.base_name}__{self.variant}"

    @property
    def id(self) -> str:
        return f"{self.base_name}__{self.variant}__{self.model}"
