import json
import random
from typing import Any

from deception_detection.data.base import DialogueDataset, DialogueDataType
from deception_detection.types import Dialogue, Label, Message
from deception_detection.utils import REPO_ROOT, get_dialogue_from_yaml


class AILiarDataset(DialogueDataset):
    """
    The dataset is adapted from Appendix F of the paper "How to Catch an AI Liar",
    https://arxiv.org/abs/2309.15840

    We have several variants:

    "original_with_answers": The original dataset. Honest dialogues and deceptive dialogues have
        different instructions. The assistant's answer is the same in both cases, and is
        human-written.
        i.e. answer held constant

    "original_without_answers"

    "mc": Multiple choice. We use the same dialogues as the original dataset,
        but append multiple choice answers to the prompt. e.g.
        ```
        Choices:
        (A) <answer>
        (B) <other_answer>
        ```
        <answer> is always the "correct" answer.
        We randomise the order of the choices.

    "normal_instruction": instruction held constant and = `normal_instruction`

    "deceive_instruction": instruction held constant and = `deceive_instruction`
    """

    base_name: str = "ai_liar"
    all_variants: list[str] = [
        "original_with_answers",
        "original_without_answers",
        "mc",
        "normal_instruction",
        "deceive_instruction",
    ]
    default_variant: str = "original_with_answers"

    def _get_dialogues(self) -> DialogueDataType:
        raw_data = self._get_raw_data()
        dialogues: list[Dialogue] = []
        labels: list[Label] = []

        for raw_dialogue in raw_data:
            honest_dialogue = self._create_dialogue(raw_dialogue, deceptive=False)
            deceptive_dialogue = self._create_dialogue(raw_dialogue, deceptive=True)

            if honest_dialogue:
                dialogues.append(honest_dialogue)
                labels.append(Label.HONEST)
            if deceptive_dialogue:
                dialogues.append(deceptive_dialogue)
                labels.append(Label.DECEPTIVE)

        return dialogues, labels, None

    def _get_raw_data(self) -> list[dict[str, str]]:
        """Get the raw data from the paper"""
        raw_data_path = REPO_ROOT / "data/how_to_catch_an_ai_liar/goal_directed_lying.json"
        with open(raw_data_path) as f:
            return json.load(f)

    def _create_dialogue(self, raw_dialogue: dict[str, str], deceptive: bool) -> Dialogue | None:
        dialogue: Dialogue = []
        match self.variant:
            case "original_with_answers":
                if not raw_dialogue["include_in_with_answers"]:
                    return None
                instruction = raw_dialogue[
                    "deceive_instruction" if deceptive else "normal_instruction"
                ]
                dialogue.append(Message("system", instruction, False))
                dialogue.append(
                    Message(
                        "user", raw_dialogue["question_prefix"] + raw_dialogue["question"], False
                    )
                )
                dialogue.append(
                    Message(
                        "assistant", raw_dialogue["answer_prefix"] + raw_dialogue["answer"], True
                    )
                )
            case "original_without_answers":
                instruction = raw_dialogue[
                    "deceive_instruction" if deceptive else "normal_instruction"
                ]
                dialogue.append(Message("system", instruction, False))
                dialogue.append(
                    Message(
                        "user", raw_dialogue["question_prefix"] + raw_dialogue["question"], False
                    )
                )
                dialogue.append(Message("assistant", raw_dialogue["answer_prefix"], True))
            case "mc":
                instruction = raw_dialogue[
                    "deceive_instruction" if deceptive else "normal_instruction"
                ]
                dialogue.append(Message("system", instruction, False))
                dialogue.append(
                    Message(
                        "user", raw_dialogue["question_prefix"] + raw_dialogue["question"], False
                    )
                )
                choices = [raw_dialogue["answer"], raw_dialogue["other_answer"]]
                correct_idx = random.choice([0, 1])
                choices = [choices[correct_idx], choices[1 - correct_idx]]
                correct_choice = chr(65 + correct_idx)
                dialogue.append(
                    Message("user", f"\n\nChoices:\n(A) {choices[0]}\n(B) {choices[1]}", False)
                )
                dialogue.append(
                    Message(
                        "assistant",
                        raw_dialogue["answer_prefix"] + " Answer: (" + correct_choice + ")",
                        True,
                    )
                )
            case "normal_instruction" | "deceive_instruction":
                instruction = raw_dialogue["normal_instruction"]
                answer = raw_dialogue["other_answer" if deceptive else "answer"]
                dialogue.append(Message("system", instruction, False))
                dialogue.append(
                    Message(
                        "user", raw_dialogue["question_prefix"] + raw_dialogue["question"], False
                    )
                )
                dialogue.append(Message("assistant", raw_dialogue["answer_prefix"] + answer, True))
            case _:
                raise ValueError(f"Unknown variant {self.variant}")
        return dialogue

    @staticmethod
    def get_grading_prompt(
        input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None = None
    ) -> dict[str, str]:
        model_grader_prompt = get_dialogue_from_yaml(
            REPO_ROOT / "data/how_to_catch_an_ai_liar/model_grader_prompt.yaml"
        )

        system_prompt = next(m.content for m in model_grader_prompt if m.role == "system")
        user_prompt = next(m.content for m in model_grader_prompt if m.role == "user")

        test_sample = ""
        for message in input_dialogue:
            if message.role == "system":
                test_sample += f"<context>{message.content}</context>\n\n"
            elif message.role == "user":
                test_sample += f"<user>{message.content}</user>\n\n"
            elif message.role == "assistant":
                test_sample += f"<assistant>{message.content}</assistant>\n\n"
        test_sample += f"<assistant>{response.content}</assistant>"

        user_prompt = user_prompt.replace("{test_sample}", test_sample.strip())
        return {"system": system_prompt, "user": user_prompt}
