from typing import cast

import pandas as pd

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


class TruthfulQA(DialogueDataset):
    base_name = "truthful_qa"

    def _get_dialogues(self) -> DialogueDataType:
        raw_data = self._get_raw_data()
        dialogues = self._get_all_dialogues(raw_data)
        labels = [Label.DECEPTIVE if i % 2 == 1 else Label.HONEST for i in range(len(dialogues))]
        return dialogues, labels, None

    def _get_raw_data(self) -> pd.DataFrame:
        """
        Get the raw data from the paper
        """
        raw_data_path = REPO_ROOT / "data/truthfulqa/eval_examples.csv"
        with open(raw_data_path) as f:
            raw_data = pd.read_csv(f)  # type: ignore
        return raw_data

    def _get_all_dialogues(self, raw_data: pd.DataFrame) -> list[Dialogue]:
        """
        Get all dialogues from the raw data, alternating between truthful and deceptive
        """
        all_dialogues: list[Dialogue] = []
        for _, row in raw_data.iterrows():
            question = cast(str, row["Question"]).strip()

            # Truthful dialogue
            truthful_dialogue: Dialogue = [
                Message("user", question, False),
                Message("assistant", cast(str, row["Examples: True"][0]), True),
            ]
            all_dialogues.append(truthful_dialogue)

            # Deceptive dialogue
            deceptive_dialogue: Dialogue = [
                Message("user", question, False),
                Message("assistant", cast(str, row["Examples: False"][0]), True),
            ]
            all_dialogues.append(deceptive_dialogue)

        return all_dialogues
