"""
DatasetRepository is used to store and retrieve both prewritten and rollout datasets.
"""

import json
import os
from collections import Counter
from pathlib import Path

import fire
import pandas as pd

import deception_detection.data as data
from deception_detection.data.base import DialogueDataset
from deception_detection.log import logger
from deception_detection.types import Label
from deception_detection.utils import REPO_ROOT

ALL_DATASETS: list[type[DialogueDataset]] = [
    data.RepeHonestDishonestFacts,
    data.AILiarDataset,
    data.AIAudit,
    data.RolePlayingDataset,
    data.RepeHonestDishonestFacts,
    data.AlpacaDataset,
    data.AILiarDataset,
    data.WerewolfDataset,
    data.GTCities,
    data.GTCitiesCitiesConj,
    data.GTCitiesCitiesDisj,
    data.GTCommonClaim,
    data.GTCompanies,
    data.GTCounterfact,
    data.GTLargerThan,
    data.GTNegCities,
    data.GTNegSpEnTrans,
    data.GTSmallerThan,
    data.GTSpEnTrans,
    data.GTMixed,
    data.ISAnimals,
    data.ISCities,
    data.ISCompanies,
    data.ISElements,
    data.ISFacts,
    data.ISInventions,
    data.ProceduralFacts,
    data.DescriptiveQA,
    data.InsiderTrading,
    data.InsiderTradingDoubleDown,
    data.SandbaggingV2Dataset,
    data.HarmRefusalChoice,
    data.HarmRefusalKnowledgeReport,
    data.InsiderTradingXYZ,
    data.InstructDishonesty,
    data.Mask,
    data.SleeperAgents,
    data.ConvinceMe,
    data.AlpacaXYZ
]


def get_id(base_name: str, variant: str, model: str) -> str:
    return f"{base_name}__{variant}__{model}"


ALL_DATASETS_DICT: dict[str, type[DialogueDataset]] = {
    dataset.base_name: dataset for dataset in ALL_DATASETS
}


def get_class_from_name(class_name: str) -> type[DialogueDataset]:
    matching = [d for d in ALL_DATASETS if d.__name__ == class_name]
    if matching:
        return matching[0]
    # if no match, return the base class
    return DialogueDataset


class DatasetRepository:
    """
    A helper class for storing and retrieving datasets.

    There are two types of datasets:

    1. Prewritten: these are datasets that are completely specified by the developer.
    They are accessed by indexing into ALL_DATASETS_DICT with a key of the form
    `BASE__VARIANT__PREWRITTEN`.

    2. Rollouts: these are datasets that are sampled on policy from some model.
    Paths to the data files are stored in the repository data file with a key of the form
    `BASE__VARIANT__MODEL`.

    We use "id" to denote to keys.
    We use "partial_id" to denote the id without the model field. This specifies a particular
    variant of a dataset.
    """

    def __init__(self, storage_file: str = "dataset_paths.json"):
        self.storage_file: Path = REPO_ROOT / storage_file
        self.dataset_paths: dict[str, str] = self._load_dataset_paths()

    def get(
        self,
        partial_id: str,
        model: str = "prewritten",
        trim_reasoning: bool = True,
        shuffle_upon_init: bool = True,
        max_llama_token_length: int | None = None,
    ) -> DialogueDataset:
        """
        Get a dataset by its partial_id.
        The model is optional, and if not specified, "prewritten" is used.
        """
        dataset_id: str = self._generate_id(partial_id, model)
        if model == "prewritten":
            base, variant = self._decode_partial_id(partial_id)
            dataset_class = ALL_DATASETS_DICT[base]
            if dataset_class is DialogueDataset:
                raise KeyError(
                    f"Prewritten dataset not found: {dataset_id}, have you added it to the "
                    "repository?"
                )
            return dataset_class(
                variant=variant,
                trim_reasoning=trim_reasoning,
                shuffle_upon_init=shuffle_upon_init,
                max_llama_token_length=max_llama_token_length,
            )

        # otherwise, it's a rollout dataset
        if dataset_id in self.dataset_paths:
            path: str = self.dataset_paths[dataset_id]
            # load the dataset class name from the file
            try:
                with open(REPO_ROOT / path) as f:
                    data = json.load(f)
            except Exception as e:
                logger.error(f"Error loading dataset {dataset_id} from {path}: {e}")
                raise e
            dataset_class_name = data.get("dataset_class", "DialogueDataset")
            base_name = data.get("base_name", "")
            variant = data.get("variant", "")
            if base_name == "" or variant == "":
                logger.warning(f"No base_name or variant found for dataset {dataset_id}")
            rollouts = data.get("rollouts", [])
            return get_class_from_name(dataset_class_name).from_rollouts(
                base_name,
                variant,
                model,
                rollouts,
                trim_reasoning,
                shuffle_upon_init,
                max_llama_token_length=max_llama_token_length,
            )

        else:
            raise KeyError(f"Dataset not found: {dataset_id}")

    def add(self, partial_id: str, model: str, path: str) -> None:
        """
        Add a dataset to the repository.
        The model is optional, and if not specified, "prewritten" is used.
        The path is the path to the data file.
        The partial_id specifies the base dataset and variant.
        """
        dataset_id: str = self._generate_id(partial_id, model)
        self.dataset_paths[dataset_id] = path
        self._save_dataset_paths()

    def check_dataset_name(self, partial_id: str, model: str = "prewritten") -> str:
        """
        Check if a dataset name is valid. If so, return the partial id.
        If not, raise a ValueError.

        Used to validate experiment configurations.
        """

        if "upper_bound_combined" in partial_id:
            return partial_id
        
        # Check if it's a prewritten dataset
        if model == "prewritten":
            base, variant = self._decode_partial_id(partial_id)
            if base in ALL_DATASETS_DICT:
                if variant in ALL_DATASETS_DICT[base].all_variants:
                    return partial_id
                else:
                    raise ValueError(f"Invalid variant for dataset {base}: {variant}")
            else:
                raise ValueError(f"Invalid prewritten dataset: {partial_id}")

        # Check if it's in the repository
        dataset_id: str = self._generate_id(partial_id, model)
        if dataset_id in self.dataset_paths:
            return partial_id
        else:
            raise ValueError(f"Dataset not found: {partial_id} with model {model}")

    def _load_dataset_paths(self) -> dict[str, str]:
        if os.path.exists(self.storage_file):
            with open(self.storage_file) as f:
                return json.load(f)
        logger.warning(f"Dataset paths file not found: {self.storage_file}")
        return {}

    def _save_dataset_paths(self) -> None:
        with open(self.storage_file, "w") as f:
            json.dump(self.dataset_paths, f, indent=2)

    def _generate_id(self, partial_id: str, model: str) -> str:
        return f"{partial_id}__{model}"

    def _decode_id(self, dataset_id: str) -> tuple[str, str, str]:
        base, variant, model = dataset_id.split("__")
        return base, variant, model

    def _decode_partial_id(self, dataset_id: str) -> tuple[str, str]:
        base, variant = dataset_id.split("__")
        return base, variant

    def print_dataset_info(self) -> None:
        """
        Print out information about all datasets in the repository.
        """
        dataset_info = []

        for base, dataset_class in ALL_DATASETS_DICT.items():
            for variant in dataset_class.all_variants:
                try:
                    dataset = dataset_class()
                    info = self._get_dataset_info(dataset, base, variant, "prewritten")
                    dataset_info.append(info)
                except Exception as e:
                    print(f"Error processing dataset {base}__{variant}: {str(e)}")

        for path in self.dataset_paths.values():
            with open(path) as f:
                data = json.load(f)
            dataset_class_name = data.get("dataset_class", "DialogueDataset")
            base_name = data.get("base_name", "")
            variant = data.get("variant", "")
            rollouts = data.get("rollouts", [])
            model = data.get("model", "")
            dataset = get_class_from_name(dataset_class_name).from_rollouts(
                base_name, variant, model, rollouts
            )
            info = self._get_dataset_info(dataset, base_name, variant, model)
            dataset_info.append(info)

        df = pd.DataFrame(dataset_info)
        df = df.sort_values(["Model", "Name"])

        print(df.to_string(index=False))

    def _get_dataset_info(
        self, dataset: DialogueDataset, base: str, variant: str, model: str
    ) -> dict[str, str | int | None]:
        name = f"{base}__{variant}"

        length = len(dataset)
        label_counts = Counter(dataset.labels)

        return {
            "Name": name,
            "Model": model,
            "Length": length,
            "Num Deceptive": label_counts[Label.DECEPTIVE],
            "Num Ambiguous": label_counts[Label.AMBIGUOUS],
            "Num Honest": label_counts[Label.HONEST],
        }


def print_dataset_info():
    repo = DatasetRepository()
    repo.print_dataset_info()


if __name__ == "__main__":
    fire.Fire(print_dataset_info)
