from __future__ import annotations

from collections.abc import Collection, Mapping
from typing import TYPE_CHECKING, Any, TypeAlias

import numpy as np
from navground import sim
from pettingzoo.utils.env import ParallelEnv

from ..config import (ActionConfig, ControlActionConfig,
                      DefaultObservationConfig, GroupConfig, ObservationConfig,
                      StateConfig)
from ..env import BaseEnv, NavgroundEnv
from ..indices import Indices, IndicesLike
from ..internal.base_env import NavgroundBaseEnv, ResetReturn, StepReturn
from ..rewards import NullReward
from ..types import (Action, Array, Bounds, Observation, Reward, SensorLike,
                     SensorSequenceLike, TerminationCondition)

if TYPE_CHECKING:
    import gymnasium as gym

BaseParallelEnv: TypeAlias = ParallelEnv[int, Observation, Action]


class MultiAgentNavgroundEnv(NavgroundBaseEnv, BaseParallelEnv):
    """
    This class describes an environment that uses
    a :py:class:`navground.sim.Scenario` to
    generate and then simulate a :py:class:`navground.sim.World`.

    Actions and observations relates to one or more selected
    :py:class:`navground.sim.Agent`.

    We provide convenience functions to initialize the class
    :py:func:`parallel_env` and :py:func:`shared_parallel_env`
    to create the environment.
    """

    metadata = {
        "name": "navground_ma",
        "render_modes": ["human", "rgb_array"],
        'render_fps': 30
    }
    MAX_SEED = 2**31 - 1

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        NavgroundBaseEnv.__init__(self, *args, **kwargs)
        self.possible_agents = list(self._possible_agents.keys())
        self.agents = list(self._agents.keys())

    def reset(self,
              seed: int | None = None,
              options: dict[str, Any] | None = None) -> ResetReturn:
        """
        Conforms to :py:meth:`pettingzoo.utils.env.ParallelEnv.reset`.

        It samples a new world from a scenario, runs one dry simulation step
        using :py:meth:`navground.sim.World.update_dry`.
        Then, it converts the agents states to observations.
        and the commands they would actuate to actions, which it includes
        in the ``info`` dictionary at key `"navground_action"``.
        """
        NavgroundBaseEnv._reset(self, seed, options)
        self.agents = list(self._agents.keys())
        return self.get_observations(), self.get_infos()

    def step(self, actions: dict[int, Action]) -> StepReturn:
        """
        Conforms to :py:meth:`pettingzoo.utils.env.ParallelEnv.step`.

        It converts the actions to commands that the navground agents actuate.
        Then, it updates the world for one step,
        calling :py:meth:`navground.sim.World.update`.
        Finally, it converts the agents states to observations,
        the commands they would actuate to actions, which it includes
        in the ``info`` dictionary at key `"navground_action"`, and computes rewards.

        Termination for individual agents is set when they complete the task,
        exit the boundary, or get stuck.
        Truncation for all agents is set when the maximal duration has passed.
        """
        obs, rew, term, trunc, info = NavgroundBaseEnv._step(self, actions)
        self.agents = [
            index for index in self.agents
            if not (trunc.get(index, False) or term.get(index, False))
        ]
        return obs, rew, term, trunc, info

    def observation_space(self, agent: int) -> gym.Space[Any]:
        return self._observation_space[agent]

    def action_space(self, agent: int) -> gym.Space[Any]:
        return self._action_space[agent]

    def state(self) -> Array:
        vs = self.get_state()
        if vs is None:
            vs = np.empty((0, ), dtype=np.float32)
        return vs


def parallel_env(
    scenario: sim.Scenario | str | dict[str, Any] | None = None,
    groups: Collection[GroupConfig] = tuple(),
    max_number_of_agents: int | None = None,
    time_step: float = 0.1,
    max_duration: float = -1.0,
    terminate_if_idle: bool = True,
    bounds: Bounds | None = None,
    truncate_outside_bounds: bool = False,
    render_mode: str | None = None,
    render_kwargs: Mapping[str, Any] = {},
    realtime_factor: float = 1.0,
    stuck_timeout: float = -1,
    wait: bool = False,
    truncate_fast: bool = False,
    state: StateConfig | None = None,
    include_action: bool = True,
    include_success: bool = True,
    init_success: bool | None = None,
    intermediate_success: bool = False,
) -> MultiAgentNavgroundEnv:
    """
    Create a multi-agent PettingZoo environment that uses
    a :py:class:`navground.sim.Scenario` to
    generate and then simulate a :py:class:`navground.sim.World`.


    >>> from navground.learning.parallel_env import parallel_env
    >>> from navground.learning import GroupConfig
    >>> from navground import sim
    >>>
    >>> scenario = sim.load_scenario(...)
    >>> groups = [GroupConfig(...), ...]
    >>> env = parallel_env(scenario=scenario, groups=groups)


    :param scenario: The scenario to initialize all simulated worlds.
                     If a :py:class:`str`, it will be interpreted as the YAML
                     representation of a scenario.
                     If a :py:class:`dict`, it will be dumped to YAML and
                     then treated as a :py:class:`str`.

    :param groups: The configuration of the agents controlled by the environment.
                   All other agents are controlled solely by navground.

    :param max_number_of_agents: The maximal number of agents that we will expose.
                                 It needs to be specified only for scenarios
                                 that generate world with a variable number of agents.

    :param time_step: The simulation time step applied
                      at every :py:meth:`MultiAgentNavgroundEnv.step`.

    :param max_duration: If positive, it will signal a truncation after this simulated time.

    :param terminate_if_idle: Whether to terminate when an agent is idle

    :param truncate_outside_bounds: Whether to truncate when an agent exit the bounds

    :param bounds: The area to render and a fence for truncating processes when agents exit it.

    :param render_mode: The render mode.
                        If `"human"`, it renders a simulation in real time via
                        websockets (see :py:class:`navground.sim.ui.WebUI`).
                        If `"rgb_array"`, it uses
                        :py:func:`navground.sim.ui.render.image_for_world`
                        to render the world on demand.

    :param render_kwargs: Arguments passed to :py:func:`navground.sim.ui.render.image_for_world`

    :param realtime_factor: a realtime factor for `render_mode="human"`: larger values
                            speed up the simulation.

    :param stuck_timeout: The time to wait before considering an agent stuck.

    :param wait: Whether to signal termination/truncation only when
                 all agents have terminated/truncated.

    :param truncate_fast: Whether to signal truncation for all agents
                          as soon as one agent truncates.

    :param state: An optional global state configuration

    :param include_action: Whether to include field "navground_action" in the info

    :param include_success: Whether to include field "is_success" in the info

    :param init_success: The default value of success (valid until a termination condition is met)

    :param intermediate_success: Whether to include "is_success" in the info at intermediate
                                 steps (vs only at termination).

    :returns:   The multi agent navground environment.
    """
    return MultiAgentNavgroundEnv(
        groups=groups,
        scenario=scenario,
        max_number_of_agents=max_number_of_agents,
        time_step=time_step,
        max_duration=max_duration,
        terminate_if_idle=terminate_if_idle,
        bounds=bounds,
        truncate_outside_bounds=truncate_outside_bounds,
        render_mode=render_mode,
        render_kwargs=render_kwargs,
        realtime_factor=realtime_factor,
        stuck_timeout=stuck_timeout,
        wait=wait,
        truncate_fast=truncate_fast,
        include_action=include_action,
        include_success=include_success,
        init_success=init_success,
        intermediate_success=intermediate_success,
        state=state)


def shared_parallel_env(
        scenario: sim.Scenario | str | dict[str, Any] | None = None,
        max_number_of_agents: int | None = None,
        indices: IndicesLike = Indices.all(),
        sensor: SensorLike | None = None,
        sensors: SensorSequenceLike = tuple(),
        action: ActionConfig = ControlActionConfig(),
        observation: ObservationConfig = DefaultObservationConfig(),
        reward: Reward | None = None,
        time_step: float = 0.1,
        max_duration: float = -1.0,
        terminate_if_idle: bool = True,
        bounds: Bounds | None = None,
        truncate_outside_bounds: bool = False,
        render_mode: str | None = None,
        render_kwargs: Mapping[str, Any] = {},
        realtime_factor: float = 1.0,
        stuck_timeout: float = -1,
        color: str = '',
        tag: str = '',
        wait: bool = False,
        truncate_fast: bool = False,
        terminate_on_success: bool = True,
        terminate_on_failure: bool = True,
        success_condition: TerminationCondition | None = None,
        failure_condition: TerminationCondition | None = None,
        include_action: bool = True,
        include_success: bool = True,
        init_success: bool | None = None,
        intermediate_success: bool = False,
        state: StateConfig | None = None) -> MultiAgentNavgroundEnv:
    """
    Create a multi-agent PettingZoo environment that uses
    a :py:class:`navground.sim.Scenario` to
    generate and then simulate a :py:class:`navground.sim.World`.
    All controlled agents share the same configuration.


    >>> from navground.learning.parallel_env import shared_parallel_env
    >>> from navground import sim
    >>>
    >>> scenario = sim.load_scenario(...)
    >>> env = shared_parallel_env(scenario=scenario, ...)

    :param scenario: The scenario to initialize all simulated worlds.
                     If a :py:class:`str`, it will be interpreted as the YAML
                     representation of a scenario.
                     If a :py:class:`dict`, it will be dumped to YAML and
                     then treated as a :py:class:`str`.

    :param max_number_of_agents: The maximal number of agents that we will expose.
                                 It needs to be specified only for scenarios
                                 that generate world with a variable number of agents.

    :param indices: The world indices of the agent to control.
                    All other agents are controlled solely by navground.

    :param sensor: An optional sensor that will be added to ``sensors``.

    :param sensors: A sequence of sensor to generate observations for the agents
                    or its YAML representation. If
                    Items of class :py:class:`str` will be interpreted as the YAML
                    representation of a sensor.
                    Items of class :py:class:`dict` will be dumped to YAML and
                    then treated as a :py:class:`str`.
                    If empty, it will use the agents' own sensors.

    :param action: The configuration of the action space to use.

    :param observation: The configuration of the observation space to use.

    :param reward: The reward function to use. If none, it will default to constant zeros.

    :param time_step: The simulation time step applied
                      at every :py:meth:`MultiAgentNavgroundEnv.step`.

    :param max_duration: If positive, it will signal a truncation after this simulated time.

    :param terminate_if_idle: Whether to terminate when an agent is idle

    :param truncate_outside_bounds: Whether to truncate when an agent exit the bounds

    :param bounds: The area to render and a fence for truncating processes when agents exit it.

    :param render_mode: The render mode.
                        If `"human"`, it renders a simulation in real time via
                        websockets (see :py:class:`navground.sim.ui.WebUI`).
                        If `"rgb_array"`, it uses
                        :py:func:`navground.sim.ui.render.image_for_world`
                        to render the world on demand.

    :param render_kwargs: Arguments passed to :py:func:`navground.sim.ui.render.image_for_world`

    :param realtime_factor: a realtime factor for `render_mode="human"`: larger values
                            speed up the simulation.

    :param stuck_timeout: The time to wait before considering an agent stuck.

    :param color: An optional color of the agents (only used for displaying)

    :param tag: An optional tag to be added to the agents (only used as metadata)

    :param wait: Whether to signal termination/truncation only when
                 all agents have terminated/truncated.

    :param truncate_fast: Whether to signal truncation for all agents
                          as soon as one agent truncates.

    :param state: An optional global state configuration

    :param terminate_on_success: Whether to terminate the episode on success.

    :param terminate_on_failure: Whether to terminate the episode on failure.

    :param success_condition: An optional success criteria

    :param failure_condition: An optional failure criteria

    :param include_action: Whether to include field "navground_action" in the info

    :param include_success: Whether to include field "is_success" in the info

    :param init_success: The default value of success (valid until a termination condition is met)

    :param intermediate_success: Whether to include "is_success" in the info at intermediate
                                 steps (vs only at termination).



    :returns:   The multi agent navground environment.
    """
    if reward is None:
        reward = NullReward()
    group = GroupConfig(indices=Indices(indices),
                        sensor=sensor,
                        sensors=sensors,
                        reward=reward,
                        action=action,
                        observation=observation,
                        color=color,
                        tag=tag,
                        terminate_on_success=terminate_on_success,
                        terminate_on_failure=terminate_on_failure,
                        success_condition=success_condition,
                        failure_condition=failure_condition)
    return MultiAgentNavgroundEnv(
        groups=[group],
        scenario=scenario,
        time_step=time_step,
        max_duration=max_duration,
        terminate_if_idle=terminate_if_idle,
        bounds=bounds,
        truncate_outside_bounds=truncate_outside_bounds,
        render_mode=render_mode,
        render_kwargs=render_kwargs,
        realtime_factor=realtime_factor,
        stuck_timeout=stuck_timeout,
        wait=wait,
        truncate_fast=truncate_fast,
        include_action=include_action,
        include_success=include_success,
        init_success=init_success,
        intermediate_success=intermediate_success,
        state=state)


def make_shared_parallel_env_with_env(
    env: BaseEnv,
    max_number_of_agents: int | None = None,
    indices: IndicesLike = Indices.all()
) -> MultiAgentNavgroundEnv:
    """
    Creates a shared parallel environment
    using the configuration of the agent exposed
    in a (navground) single-agent environment.


    :param      env:                   The environment.
    :param      max_number_of_agents:  The maximal number of agents that we will expose.
        It needs to be specified only for scenarios
        that generate world with a variable number of agents.
    :param      indices:         The world indices of the agent to control.
        All other agents are controlled solely by navground.

    :raises TypeError: If ``env.unwrapped`` is not a subclass of :py:class:`.env.NavgroundEnv`

    """
    if not isinstance(env.unwrapped, NavgroundEnv):
        raise TypeError("Environment is not a NavgroundEnv")

    kwargs = dict(**env.unwrapped._spec)
    kwargs.pop('agent_index')
    kwargs.pop('reward')
    kwargs.pop('action')
    kwargs.pop('observation')
    kwargs.pop('sensor')
    kwargs.pop('sensors')
    kwargs['max_number_of_agents'] = max_number_of_agents
    group = list(env.unwrapped.groups_config)[0]
    group.indices = Indices(indices)
    return MultiAgentNavgroundEnv(groups=[group], **kwargs)
