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

import functools
from typing import Any, Dict, Optional

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

import chatarena
from chatarena.arena import Arena

import string

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 OpenSpiel 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 == True:
            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 == True:
            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()



