import copy
import warnings
from contextlib import contextmanager
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple

import gymnasium as gym
import torch
from accelerate import Accelerator
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from transformers.tokenization_utils_base import BatchEncoding


class HuggingFaceGym(gym.Env):
    """Class to convert HuggingFace datasets into Gymnasium style environment.

    :param dataset_name: Dataset name to be loaded from HuggingFace datasets.
    :type dataset_name: str
    :param tokenizer: Tokenizer to be used for encoding and decoding the prompts.
    :type tokenizer: AutoTokenizer
    :param reward_fn: Reward function for evaluating completions.
    :type reward_fn: Callable[..., float]
    :param data_batch_size_per_gpu: DataLoader batch size, defaults to 8
    :type data_batch_size_per_gpu: int, optional
    :param custom_collate_fn: Custom collate function to be used for creating the batch, defaults to None
    :type custom_collate_fn: Callable, optional
    :param accelerator: Accelerator to be used for training, defaults to None
    :type accelerator: Accelerator, optional
    """

    def __init__(
        self,
        train_dataset: str,
        test_dataset,
        tokenizer: AutoTokenizer,
        reward_fn: Callable[[str, str, str], float],
        apply_chat_template_fn: Callable[[str, str, AutoTokenizer], BatchEncoding],
        data_batch_size_per_gpu: int = 8,
        custom_collate_fn: Optional[Callable] = None,
        accelerator: Optional[Accelerator] = None,
    ) -> None:
        assert {"question", "answer"}.issubset(
            set(train_dataset.features.keys())
        ), "Train dataset must contain 'question' and 'answer' features."
        assert {"question", "answer"}.issubset(
            set(test_dataset.features.keys())
        ), "Train dataset must contain 'question' and 'answer' features."
        self.name = train_dataset.info.dataset_name
        self.reward_fn = reward_fn
        self.tokenizer = tokenizer
        self.data_batch_size_per_gpu = data_batch_size_per_gpu
        if custom_collate_fn is None:
            custom_collate_fn = HuggingFaceGym.create_collate_fn(
                tokenizer, apply_chat_template_fn
            )
        dataloader_kwargs = {"collate_fn": custom_collate_fn}
        self.train_dataloader = DataLoader(
            train_dataset,
            batch_size=data_batch_size_per_gpu,
            shuffle=True,
            **dataloader_kwargs,
        )
        self.test_dataloader = DataLoader(
            test_dataset,
            batch_size=data_batch_size_per_gpu,
            shuffle=False,
            **dataloader_kwargs,
        )
        self.dataset_size = {
            "train": len(train_dataset),
            "test": len(test_dataset),
        }
        self.accelerator = accelerator
        if self.accelerator is not None:
            self.train_dataloader = self.accelerator.prepare(self.train_dataloader)
            self.test_dataloader = self.accelerator.prepare(self.test_dataloader)
        self.train_dataloader_iter = iter(self.train_dataloader)
        self.test_dataloader_iter = iter(self.test_dataloader)
        self.apply_chat_template_fn = apply_chat_template_fn
        self.dataloader = self.train_dataloader_iter
        self.reset_called = False
        self.observation_space = gym.spaces.Box(low=0, high=tokenizer.vocab_size - 1)
        self.action_space = gym.spaces.Box(
            low=0,
            high=tokenizer.vocab_size - 1,
        )
        self.evaluation_mode = False
        self.num_dataset_passes = 0

    def step(
        self, completions: torch.Tensor
    ) -> Tuple[List[BatchEncoding], torch.Tensor]:
        """Take a step in the HuggingFaceGym environment, calculate rewards from completions generated from previous prompt and provide new batch
        of prompts.

        :param completions: Completion IDs generated by the agent.
        :type completions: torch.Tensor
        :return: New tokenized prompts and an information dictionary.
        :rtype: Tuple[List[BatchEncoding], torch.Tensor]
        """
        self.reset_called = False
        rewards = self._decode_and_evaluate(completions)
        new_tokenized_prompts = self._get_next_batch()
        self.last_tokenized_prompts = new_tokenized_prompts
        return new_tokenized_prompts, rewards

    def reset(
        self, reset_dataloaders: bool = False
    ) -> Tuple[List[BatchEncoding], Dict[str, Any]]:
        """Reset the environment and get the next batch of tokenized prompts.

        :param reset_dataloaders: Whether to reset the dataloaders, defaults to False
        :type reset_dataloaders: bool, optional
        :return: New tokenized prompts
        :rtype: Tuple[List[BatchEncoding], Dict[str, Any]]
        """
        if reset_dataloaders:
            self._reset_dataloaders()
            warnings.warn(
                "env.reset() called with reset_dataloaders=True, this will reset the dataloaders to the beginning of the dataset, proceed with caution."
            )
        if self.reset_called:
            warnings.warn(
                "env.reset() called more than once sequentially, it should typically follow with env.step()."
            )
        self.reset_called = True
        new_tokenized_prompts = self._get_next_batch()
        self.last_tokenized_prompts = new_tokenized_prompts
        return new_tokenized_prompts

    def _decode_and_evaluate(self, completions: List[torch.Tensor]) -> torch.Tensor:
        """Decode the completions and evaluate the rewards.

        :param completions: Completion IDs generated by the agent.
        :type completions: List[torch.Tensor]
        :return: Rewards for the completions.
        :rtype: torch.Tensor
        """
        # This is for a batch of completions (prompt_batch x group_size), List of tensors of length batch size, each tensor is a group of answers
        total_rewards = []
        for idx, (group_completion, answer, question) in enumerate(
            zip(completions, self.answers, self.questions)
        ):  # Vectorize this in the future
            decoded_group_completion = self.tokenizer.batch_decode(
                group_completion[
                    :, self.last_tokenized_prompts[idx]["input_ids"].shape[1] :
                ],
                skip_special_tokens=True,
            )
            rewards = [
                self.reward_fn(completion, answer, question)
                for completion in decoded_group_completion
            ]
            total_rewards.append(rewards)
        return torch.tensor(total_rewards)

    def _get_next_batch(self) -> List[BatchEncoding]:
        """Get the next batch of tokenized prompts."""
        try:
            batch = next(self.dataloader)
            self.questions = batch["question"]
            self.answers = batch["answer"]
            tokenized_prompts = batch["tokenized_prompts"]
        except StopIteration:
            self._reset_dataloaders(
                reset_train=not self.evaluation_mode,
                reset_test=self.evaluation_mode,
            )
            self.num_dataset_passes += 1
            return self._get_next_batch()
        return tokenized_prompts

    @contextmanager
    def eval_mode(self) -> Generator[None, None, None]:
        """Context manager to switch to evaluation mode."""
        self.dataloader = self.test_dataloader_iter
        self.evaluation_mode = True
        last_tokenized_prompts = copy.deepcopy(self.last_tokenized_prompts)
        try:
            yield
        finally:
            self.dataloader = self.train_dataloader_iter
            self.evaluation_mode = False
            self.last_tokenized_prompts = last_tokenized_prompts

    def __len__(self):
        """Return the length of the dataset."""
        if self.evaluation_mode:
            return len(self.test_dataloader.dataset)
        return len(self.train_dataloader.dataset)

    def _reset_dataloaders(self, reset_train: bool = True, reset_test: bool = True):
        """Reset the dataloaders to the beginning of the dataset.

        :param reset_train: Whether to reset the train dataloader, defaults to True
        :type reset_train: bool, optional
        :param reset_test: Whether to reset the test dataloader, defaults to True
        :type reset_test: bool, optional
        """
        if reset_train:
            self.train_dataloader_iter = iter(self.train_dataloader)
        if reset_test:
            self.test_dataloader_iter = iter(self.test_dataloader)
        self.dataloader = (
            self.test_dataloader_iter
            if self.evaluation_mode
            else self.train_dataloader_iter
        )

    @staticmethod
    def create_collate_fn(
        tokenizer: AutoTokenizer,
        apply_chat_template_fn: Callable[[str, str, AutoTokenizer], BatchEncoding],
    ) -> Callable[[List[Dict[str, Any]]], Dict[str, Any]]:
        """
        Create a collate function that applies the chat template to the batch of questions and answers.

        :param tokenizer: Tokenizer to be used for encoding and decoding the prompts.
        :type tokenizer: AutoTokenizer
        :param apply_chat_template_fn: Function to apply the chat template to the batch of questions and answers.
        :type apply_chat_template_fn: Callable[[str, str, AutoTokenizer], BatchEncoding]
        :return: Collate function that applies the chat template to the batch of questions and answers.
        :rtype: Callable[[List[Dict[str, Any]]], Dict[str, Any]]
        """

        def collate_fn(batch):

            questions = [item["question"] for item in batch]
            answers = [item["answer"] for item in batch]

            # Apply chat template to all samples
            tokenized_prompts = [
                apply_chat_template_fn(q, a, tokenizer)
                for q, a in zip(questions, answers)
            ]

            return {
                "question": questions,
                "answer": answers,
                "tokenized_prompts": tokenized_prompts,  # Keep individual tokenized prompts
            }

        return collate_fn


class DummyOptimizer:
    """
    Placeholder optimizer class to pass to the OptimizerWrapper when the optimizer is defined in the deepspeed config.
    """

    def __init__(self, params: List[torch.Tensor], lr: float, **kwargs) -> None:
        """
        Sentinel class to use for the optimizer when the optimizer is defined in the deepspeed config.

        :param params: Parameters to optimize.
        :type params: List[torch.Tensor]
        :param lr: Learning rate.
        :type lr: float
        """
        pass

    def step(self, closure=None):
        raise RuntimeError(
            "DummyOptimizer is a placeholder optimizer and should not be used."
            "Please ensure you are calling accelerator.prepare() on the optimizer."
        )

    def zero_grad(self):
        raise RuntimeError(
            "DummyOptimizer is a placeholder optimizer and should not be used."
            "Please ensure you are calling accelerator.prepare() on the optimizer."
        )

    def state_dict(self):
        raise RuntimeError(
            "DummyOptimizer is a placeholder optimizer and should not be used."
            "Please ensure you are calling accelerator.prepare() on the optimizer."
        )

    def load_state_dict(self, state_dict):
        raise RuntimeError(
            "DummyOptimizer is a placeholder optimizer and should not be used."
            "Please ensure you are calling accelerator.prepare() on the optimizer."
        )
