"""Wrapper to convert a ChatArena environment into a PettingZoo compatible environment."""
from __future__ import annotations

import functools
import string

import pettingzoo
from gymnasium import spaces
from pettingzoo.utils.env import AgentID, ObsType

import chatarena
from chatarena.arena import Arena

CHAR_SET = string.printable


class PettingZooCompatibilityV0(pettingzoo.AECEnv):
    """This compatibility wrapper converts a ChatArena environment into a PettingZoo environment.

    ChatArena (or Chat Arena) is a Multi-Agent Language Game Environments for LLMs. The goal is to develop communication and collaboration capabilities of AIs.
    """

    metadata = {
        "render_modes": ["human"],
        "name": "PettingZooCompatibilityV0",
        "is_parallelizable": False,
    }

    def __init__(
        self,
        env: chatarena.arena.Arena | None = None,
        arena_name: str | None = None,
        string_observation: bool | None = True,
        max_turns: int | None = 25,
        render_mode: str | None = None,
    ):
        """Wrapper to convert a ChatArena environment into a PettingZoo environment.

        Args:
            env (chatarena.arena.Arena): chatarena arena to wrap
            arena_name (Optional[str]): chatarena environment to load from file (e.g., "examples/chameleon.json")
            max_turns (Optional[int]): maximum number of turns before environment truncates
            render_mode (Optional[str]): rendering mode
        """
        super().__init__()

        if env is not None:
            self._env = env
        elif arena_name is not None:
            self._env = Arena.from_config(arena_name)
        else:
            raise ValueError(
                "Arena not specified, please us env or arena_name arguments."
            )

        self._env.reset()  # this resets the underlying arena as well as each player

        self.possible_agents = list(self._env.name_to_player.keys())
        self.name_to_player_mapping = self._env.name_to_player

        self.string_observation = string_observation
        self.max_turns = max_turns
        self.render_mode = render_mode

        self.terminations = {}
        self.truncations = {}
        self.rewards = {}
        self.infos = {a: {} for a in self.possible_agents}

    @functools.lru_cache(maxsize=None)
    def observation_space(self, agent: AgentID):
        """Observation_space.

        We get the observation space from the underlying environment.
        Args:
            agent (AgentID): agent
        """
        # TODO: finalize obs space (dicts with agent name may not be necessary)
        observation_space = spaces.Dict(
            {
                agent: spaces.Text(max_length=256, min_length=0, charset=CHAR_SET)
                for agent in self.possible_agents
            }
        )
        return observation_space

    @functools.lru_cache(maxsize=None)
    def action_space(self, agent: AgentID):
        """Action_space.

        Get the action space from the underlying environment.

        Args:
            agent (AgentID): agent

        Returns:
            space
        """
        # TODO: finalize action space (this enables agents to send messages to specific other players)
        action_space = spaces.Dict(
            {
                agent: spaces.Text(max_length=256, min_length=0, charset=CHAR_SET)
                for agent in self.possible_agents
            }
        )
        return action_space

    def render(self):
        """Render.

        Print the current game state.
        """
        if not hasattr(self, "initial_timestep"):
            raise UserWarning(
                "You must reset the environment using reset() before calling render()."
            )

        self._env.environment.print()
        pass

    def observe(self, agent: AgentID) -> ObsType:
        """Observe.

        Args:
            agent (AgentID): agent (e.g., "Player 1")

        Returns:
            observation
        """
        messages = self._env.environment.get_observation(
            agent
        )  # this will only return the messages this agent can see
        if len(messages) > 0:
            self.current_turn = messages[-1].turn
        else:
            self.current_turn = 0
        new_messages = [
            m for m in messages if m.turn == self.current_turn
        ]  # we only send the current timestep messages

        # string observation
        if self.string_observation:
            observation = ""
            for m in new_messages:
                observation += f"{m.agent_name}: {m.content}"

        # dict observation
        else:
            observation = {m.agent_name: m.content for m in new_messages}

        self.infos[agent]["obs_dict"] = {m.agent_name: m.content for m in new_messages}

        return observation

    def close(self):
        """Close."""
        pass

    def _unravel_timestep(self, timestep: chatarena.arena.TimeStep):
        # get observation
        messages = timestep.observation
        if len(messages) > 0:
            self.current_turn = messages[-1].turn
        else:
            self.current_turn = 0
        new_messages = [
            m for m in messages if m.turn == self.current_turn
        ]  # we only send the current timestep messages

        # string observation
        if self.string_observation:
            observation = ""
            for m in new_messages:
                observation += f"{m.agent_name}: {m.content}"

        # dict observation
        else:
            observation = {m.agent_name: m.content for m in new_messages}

        # get rewards
        rewards = timestep.reward

        # get termination
        termination = timestep.terminal

        # get truncation
        truncation = self.current_turn > self.max_turns

        # get info
        player_idx = self.possible_agents.index(self.agent_selection)
        player_obj = self._env.players[player_idx]
        info = {
            "turn": self.current_turn,
            "global_prompt": player_obj.global_prompt,
            "agent_desc": player_obj.role_desc,
        }

        return observation, rewards, termination, truncation, info

    def reset(
        self,
        return_info: bool | None = False,
        seed: int | None = None,
        options: dict | None = None,
    ):
        """Reset.

        Args:
            seed (Optional[int]): seed
            options (Optional[Dict]): options
        """

        if seed is not None:
            print("WARNING: seeding is not supported for LLMs.")

        # reset the chat arena environment
        self.initial_timestep = self._env.reset()
        self.turn = 0

        # get the first player
        self.agent_selection = self._env.environment.get_next_player()

        observation, reward, termination, truncation, info = self._unravel_timestep(
            self.initial_timestep
        )

        agent = self.agent_selection
        self.rewards = reward
        self.terminations[agent] = termination
        self.truncations[agent] = truncation
        self.infos[agent] = info

        # all agents
        self.agents = self.possible_agents[:]

        # boilerplate stuff
        self._cumulative_rewards = {a: 0 for a in self.agents}
        self.rewards = self.initial_timestep.reward
        self.terminations = {a: False for a in self.agents}
        self.truncations = {a: False for a in self.agents}

    def step(self, action: str):
        """Steps.

        Steps the agent with an action.

        Args:
            action (str): action
        """
        if (
            self.terminations[self.agent_selection]
            or self.truncations[self.agent_selection]
        ):
            return self._was_dead_step(action)

        agent = self.agent_selection
        timestep = self._env.environment.step(player_name=agent, action=action)

        observation, reward, termination, truncation, info = self._unravel_timestep(
            timestep
        )

        self.rewards = reward
        self.terminations[agent] = termination
        self.truncations[agent] = truncation
        self.infos[agent] = info

        self._cumulative_rewards[agent] = 0

        self.agent_selection = self._env.environment.get_next_player()

        self._accumulate_rewards()

        if self.render_mode == "human":
            self.render()
