from __future__ import annotations

import abc
import dataclasses as dc
import math
from collections.abc import Mapping
from operator import attrgetter
from typing import TYPE_CHECKING, Any

import numpy as np
from navground import core, sim

from ..register import Registrable
from ..types import Action, Array, Observation

if TYPE_CHECKING:
    import gymnasium as gym

BoxType = type[np.floating[Any]] | type[np.integer[Any]]


@dc.dataclass(repr=False)
class DataclassConfig:
    dtype: str = ''
    """The data type"""

    def __post_init__(self) -> None:
        if self.dtype:
            self._dtype: np.dtype[Any] | None = np.dtype(self.dtype)
            self.dtype = self._dtype.name
        else:
            self._dtype = None

    # Same as dtype accepted in `gym.Box(dtype=...)`
    @property
    def box_type(self) -> BoxType:
        if self._dtype:
            return self._dtype.type
        return core.FloatType

    def __repr__(self) -> str:
        nodef_f_vals = ((f.name, attrgetter(f.name)(self))
                        for f in dc.fields(self)
                        if attrgetter(f.name)(self) != f.default)
        nodef_f_repr = ", ".join(f"{name}={value}"
                                 for name, value in nodef_f_vals)
        return f"{self.__class__.__name__}({nodef_f_repr})"

    def _get_dict(self) -> dict[str, Any]:
        return dc.asdict(self)

    @classmethod
    def _make_from_dict(cls, value: Mapping[str, Any]) -> Any:
        return cls(**value)


@dc.dataclass(repr=False)
class ConfigWithKinematic(DataclassConfig):
    """
    An internal helper class that configures kinematics properties
    """

    max_speed: float = np.inf
    """The maximal speed of the agent"""
    max_angular_speed: float = np.inf
    """The maximal angular speed of the agent"""
    dof: int | None = None
    """The number of degrees of freedom of the agent"""

    def configure_kinematics(self, behavior: core.Behavior) -> None:
        """
        Configure the fields that are not yet set from the behavior's
        kinematics.

        :param      behavior:  The behavior
        """
        kinematics = behavior.kinematics
        if not kinematics:
            return
        if not math.isfinite(self.max_speed):
            self.max_speed = kinematics.max_speed
        if not math.isfinite(self.max_angular_speed):
            self.max_angular_speed = kinematics.max_angular_speed
        if self.dof is None:
            self.dof = kinematics.dof()


class ObservationConfig(abc.ABC, Registrable):
    """
    Configure how to convert the state of a
    a :py:class:`navground.sim.Agent` to observation consumed
    in a :py:class:`gymnasium.Env`.

    Observation may contain data from the agent's behavior ego state (e.g., velocity)
    and :py:attr:`navground.core.Behavior.target` (e.g. relative target direction),
    and data generated by :py:class:`navground.sim.Sensor`.
    """

    @property
    @abc.abstractmethod
    def history(self) -> int:
        """
        The size of observations queue.
        If larger than 1, recent observations will be first
        stacked and then flattened.
        """
        ...

    @abc.abstractmethod
    def configure(self, behavior: core.Behavior | None,
                  sensing_space: gym.spaces.Dict) -> None:
        """
        Complete the configuration for a target behavior is known.

        :param      behavior:  The behavior
        :param      sensing_space:  The sensing space
        """
        ...

    @abc.abstractmethod
    def is_configured(self, warn: bool = False) -> bool:
        """
        Whether is the configuration is complete.

        :returns:   True if complete, False otherwise.
        """
        ...

    # @property
    # @abc.abstractmethod
    # def asdict(self) -> dict[str, Any]:
    #     """
    #     A JSON-able representation of the configuration

    #     :returns:  A JSON-able dict
    #     """
    #     ...

    @abc.abstractmethod
    def get_observation(
            self, behavior: core.Behavior | None,
            buffers: Mapping[str, core.Buffer]) -> Observation:
        ...

    @property
    @abc.abstractmethod
    def space(self) -> gym.Space[Any]:
        ...


@dc.dataclass
class ActionConfig(abc.ABC, Registrable):
    """
    The abstract base time of actions configuration.

    Describes how the agent should convert control commands
    from navground :py:class:`navground.core.Behavior`
    to actions to be applied a :py:class:`gymnasium.Env`
    and, in the opposite direction), actions computed by a ML-policy
    to control commands to be applied in navground.
    """

    @abc.abstractmethod
    def get_cmd_from_action(self, action: Action,
                            behavior: core.Behavior | None,
                            time_step: float) -> core.Twist2:
        """
        Convert action to navground command.

        :param      action:     The action
        :param      behavior:   The behavior that is requesting the command
        :param      time_step:  The time step

        :returns:   A control command.
        """
        ...

    @abc.abstractmethod
    def get_action(self, behavior: core.Behavior, time_step: float) -> Action:
        """
        Gets the action from the command computed by a behavior.

        :param      behavior:   The behavior
        :param      time_step:  The time step

        :returns:   The action.
        """
        ...

    @property
    @abc.abstractmethod
    def space(self) -> gym.Space[Any]:
        """
        The action space

        :returns:   A valid Gymnasium space
        """
        ...

    @abc.abstractmethod
    def configure(self, behavior: core.Behavior) -> None:
        """
        Complete the configuration for a target behavior is known.

        :param      behavior:  The behavior
        """
        ...

    @abc.abstractmethod
    def is_configured(self, warn: bool = False) -> bool:
        """
        Whether is the configuration is complete.

        :returns:   True if complete, False otherwise.
        """
        ...

    # @property
    # @abc.abstractmethod
    # def asdict(self) -> dict[str, Any]:
    #     """
    #     A JSON-able representation of the configuration

    #     :returns:  A JSON-able dict
    #     """
    #     ...


class StateConfig(abc.ABC, Registrable):
    """
    Configure how to get the environment state from the current state of
    a :py:class:`navground.sim.World`
    """

    @abc.abstractmethod
    def get_state(self, world: sim.World) -> Array:
        """
          Computes the global state from the current world state

          :param      world:  The world

          :returns: an array with the global state
        """
        ...

    @abc.abstractmethod
    def get_space(self, world: sim.World) -> gym.spaces.Box:
        """
        Returns the space of the global state

        :param      world:  The world

        :returns:   The space.
        """
        ...
