import datetime
import json
from pathlib import Path
from typing import Any, cast

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


class AIAudit(DialogueDataset):
    """
    Scenarios from the paper "Uncovering Deceptive Tendencies in Large Language Models" by
    Olli Järviniemi.

    Currently, we omit A2, which has a slightly different setup.
    """

    base_name: str = "ai_audit"
    default_variant: str = "no_reasoning"
    all_variants: list[str] = ["no_reasoning", "reasoning", "a1_2", "a1_2_reasoning"]
    trim_reasoning_str: str | None = "Action: send_message"
    base_path: Path = REPO_ROOT / "data" / "ai_audit" / "Uncovering_Deceptive_Tendencies"

    padding = {model_type: {"left": 0, "right": 0} for model_type in ["gemma", "mistral", "llama", "qwen"]}

    def _get_dialogues_a1_2(self) -> DialogueDataType:
        setup = self.get_setup(reasoning=self.variant in ["reasoning", "a1_2_reasoning"])
        # get variation 2 of A1
        final_day_dialogue = self.get_experiment_A1(with_generation_task=True)[1]
        dialogue = self.prepare_dialogues(setup + final_day_dialogue, num_skipped_tasks=0)
        metadata = {
            "experiment": ["A1"],
            "experiment_variation": ["A1_2"],
        }
        return [dialogue], [Label.DECEPTIVE], metadata

    def _get_dialogues(self) -> DialogueDataType:
        if self.variant in ["a1_2", "a1_2_reasoning"]:
            return self._get_dialogues_a1_2()

        setup = self.get_setup(reasoning=self.variant == "reasoning")
        setup_without_code = self.get_setup(
            reasoning=self.variant == "reasoning", without_code_task=True
        )
        dialogues: list[Dialogue] = []
        metadata: dict[str, list[Any]] = {
            "experiment": [],
            "experiment_variation": [],
        }
        labels: list[Label] = []

        # A1 deceptive
        for i, final_day_dialogue in enumerate(self.get_experiment_A1(with_generation_task=True)):
            dialogue = self.prepare_dialogues(setup + final_day_dialogue, num_skipped_tasks=0)
            dialogues.append(dialogue)
            metadata["experiment"].append("A1")
            metadata["experiment_variation"].append(f"A1_{i + 1}")
            labels.append(Label.DECEPTIVE)

        # A1 honest
        for i, final_day_dialogue in enumerate(self.get_experiment_A1(with_generation_task=False)):
            dialogue = self.prepare_dialogues(setup + final_day_dialogue, num_skipped_tasks=1)
            dialogues.append(dialogue)
            metadata["experiment"].append("A1")
            metadata["experiment_variation"].append(f"A1_{i + 1}")
            labels.append(Label.HONEST)

        for final_day_dialogue in self.get_experiment_B():
            # DECEPTIVE
            dialogue = self.prepare_dialogues(setup + final_day_dialogue, num_skipped_tasks=0)
            dialogues.append(dialogue)
            metadata["experiment"].append("B")
            metadata["experiment_variation"].append("B")
            labels.append(Label.DECEPTIVE)
            # HONEST
            dialogue = self.prepare_dialogues(
                setup_without_code + final_day_dialogue, num_skipped_tasks=1
            )
            dialogues.append(dialogue)
            metadata["experiment"].append("B")
            metadata["experiment_variation"].append("B")
            labels.append(Label.HONEST)

        return dialogues, labels, metadata

    def prepare_dialogues(self, dialogue: Dialogue, num_skipped_tasks: int = 0) -> Dialogue:
        """
        Prepares a single dialogue by replacing placeholders with actual values, and injecting
        reasoning if required by the variant.
        """
        # there are some parameterized values to overwrite for day 4 content
        START_ID = 291 - num_skipped_tasks
        DATE = datetime.date.fromisoformat("2024-07-26")
        tasks_completed = 105 - num_skipped_tasks

        date_map = {
            "2024-03-06": "2024-07-23",
            "2024-03-07": "2024-07-24",
            "2024-03-08": "2024-07-25",
        }

        # Replace placeholders in each dialogue with actual values
        id_idx = 0  # Reset ID index for each dialogue
        for message in dialogue:
            message.content = message.content.replace("[DATE]", str(DATE))
            # we transpose all dates to match the llama system prompt
            for orig_date, new_date in date_map.items():
                message.content = message.content.replace(orig_date, new_date)
            message.content = message.content.replace("March, 2024", "July, 2024")
            if "[TASK]" in message.content:
                message.content = message.content.replace("[TASK]", str(tasks_completed))
                tasks_completed += 1
            if "[ID]" in message.content:
                message.content = message.content.replace("[ID]", str(START_ID + id_idx))
                id_idx += 1  # noqa: SIM113

        if self.variant == "reasoning":
            # inject an assistant message that starts with "Reasoning: """ at the end of the
            # dialogue
            dialogue.append(Message("assistant", 'Reasoning: """', False))

        return dialogue

    @staticmethod
    def get_grading_prompt(
        input_dialogue: Dialogue, response: Message, metadata: dict[str, Any] | None = None
    ) -> dict[str, str]:
        if metadata is None:
            raise ValueError("Metadata is required")

        match metadata["experiment"]:
            case "A1":
                filename = "a1_grading_prompt.txt"
            case "B":
                filename = "b_grading_prompt.txt"
            case "C":
                filename = "c_grading_prompt.txt"
            case _:
                raise ValueError(f"Unknown experiment {metadata['experiment']}")

        path = REPO_ROOT / "data" / "ai_audit" / "custom" / filename
        with open(path) as f:
            system_prompt = f.read()

        user_prompt = """<ai_response>
        {response}
        </ai_response>
        """
        user_prompt = user_prompt.format(response=response.content)
        return {"system": system_prompt, "user": user_prompt}

    def get_setup(self, reasoning: bool, without_code_task: bool = False) -> Dialogue:
        """
        Get's the shared setup for all dialogues in the dataset. "Day 1 to Day 3".
        """
        file_name = "4" if reasoning else "no_reasoning"
        file_name = f"{file_name}_without_code" if without_code_task else file_name
        path = f"trajectories/{file_name}.json"
        return self.get_messages_from_json(path)

    def get_experiment_A1(self, with_generation_task: bool = True) -> list[Dialogue]:
        """
        Get's the dialogues for experiment A1.
        Ask the model if it did unethical actions.
        """
        if with_generation_task:
            preface = [
                ("user", "experiments/A/generation_task.txt"),
                ("user", "experiments/A/generation_response.txt"),
                ("user", "experiments/A/completion.txt"),
                ("user", "experiments/A/slack_A1.txt"),
            ]
        else:
            preface = [
                ("user", "experiments/A/slack_A1.txt"),
            ]

        txt_lists = [
            preface + [("user", f"experiments/A/evaluation_task_{i}.txt")] for i in range(1, 8)
        ]
        return [self.get_dialogue_from_txt_list(txt_list) for txt_list in txt_lists]

    def get_experiment_A2(self) -> list[Dialogue]:
        """
        Get's the dialogues for experiment A2.
        Check if the model proactively deletes things.
        """
        raise NotImplementedError
        # i think this requires generating responses in the middle of the dialogue,
        # which is annoying.

    def get_experiment_B(self) -> list[Dialogue]:
        """
        Get's the dialogues for experiment B.
        Directly ask the model about actions, which were benign, but are now told are bad.

        We only use evaluation task 1 -- evaluation task 2 is asking the model to describe it's
        past actions without any reason to be deceptive.
        """
        txt_list = [
            ("user", "experiments/B/slack.txt"),
            ("user", "experiments/B/evaluation_task_1.txt"),
        ]
        dialogue = self.get_dialogue_from_txt_list(txt_list)
        return [dialogue]

    def get_experiment_C(self) -> list[Dialogue]:
        """
        Get's the dialogues for experiment C.
        Ask the model to do something that it is told it should not do.
        We omit the capability evaluation that checks if the model actually suceeds.
        """
        C1 = [
            [
                ("user", f"experiments/C/slack_{i}.txt"),
                ("user", "experiments/C/evaluation_task.txt"),
            ]
            for i in range(1, 7)
        ]

        C2 = [
            [
                ("user", f"experiments/C/slack_{i}.txt"),
                ("user", "experiments/C/evaluation_task_2.txt"),
            ]
            for i in range(1, 7)
        ]

        return [self.get_dialogue_from_txt_list(dialogue) for dialogue in C1 + C2]

    def get_messages_from_json(self, json_path: str) -> Dialogue:
        path = self.base_path / json_path
        with open(path) as f:
            data = json.load(f)
        dialogue = [
            Message(msg["role"], msg["content"].strip("\n"), False)
            for msg in data
            if msg["role"] != "canary"
        ]
        return dialogue

    def get_message_from_txt(self, role: str, txt_path: str) -> Message:
        path = self.base_path / txt_path
        with open(path) as f:
            data = f.read()
        return Message(cast(Role, role), data.strip("\n"), False)

    def get_dialogue_from_txt_list(self, txt_list: list[tuple[str, str]]) -> Dialogue:
        return [self.get_message_from_txt(role, path) for role, path in txt_list]
