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

REASONING_SYSTEM_PROMPT = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
    "<think> reasoning process here </think><answer> answer here </answer>"
)


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 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: 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
        dataloader_kwargs = (
            {} if custom_collate_fn is None else {"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.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.eval_mode = False

    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]]:
        if reset_dataloaders:
            self._reset_dataloaders()
        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:
        # 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 = []
        if self.eval_mode:
            decoded_completions = []
        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,
            )
            if self.eval_mode:
                decoded_completions.append(decoded_group_completion)
            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]:
        batch = next(self.dataloader)
        self.questions = batch["question"]
        self.answers = batch["answer"]
        tokenized_prompts = [
            self.apply_chat_template_fn(question, answer, self.tokenizer)
            for question, answer in zip(self.questions, self.answers)
        ]
        return tokenized_prompts

    @contextmanager
    def eval(self) -> Generator[None, None, None]:
        self.dataloader = self.test_dataloader_iter
        self.eval_mode = True
        try:
            yield
        finally:
            self.dataloader = self.train_dataloader_iter
            self.eval_mode = False

    def __len__(self):
        if self.eval_mode:
            return len(self.test_dataloader.dataset)
        return len(self.train_dataloader.dataset)

    def _reset_dataloaders(self):
        self.train_dataloader_iter = iter(self.train_dataloader)
        self.test_dataloader_iter = iter(self.test_dataloader)
