"""
Agents Connector Base Method
=============================

This module defines the base class for connector functions that allow the combination of 
agents' observations to provide a flexible configuration of the communication between agents. 
Built-in hierarchical (only two levels), fully-shared, centralized, and independent configurations 
are provided.
"""

from typing import Dict, Any, Tuple
from gymnasium.spaces import Space

class BaseConnector:
    def __init__(
        self,
        connector_fn_config: Dict[str, Any] = {}
    ):
        """
        Base class for connector functions.
        
        :param connector_fn_config: Configuration of the connector function.
        :type connector_fn_config: Dict[str, Any], optional
        """
        self.connector_fn_config = connector_fn_config
    
    def get_agent_obs_dim(
        self,
        env_config: Dict[str, Any],
        agent: str = None
    ) -> Space:
        """
        Get the agent observation dimension.

        :param env_config: Environment configuration.
        :type env_config: Dict[str, Any]
        :param agent: Agent identifier, optional.
        :type agent: str, optional
        :return: Agent observation spaces.
        :rtype: gym.Space
        """
        raise NotImplementedError("This method must be implemented in the child class.")
        
    def set_top_level_obs(
        self,
        env_config: Dict[str, Any],
        agent_states: Dict[str, Dict[str, Any]],
        dict_agents_obs: Dict[str, Any],
        infos: Dict[str, Dict[str, Any]],
        is_last_timestep: bool = False
    ) -> Tuple[Dict[str, Any], Dict[str, Dict[str, Any]], bool]:
        """
        Set the multi-agent observation.

        :param env_config: Environment configuration.
        :type env_config: Dict[str, Any]
        :param agent_states: Agent states.
        :type agent_states: Dict[str, Dict[str, Any]]
        :param dict_agents_obs: Dictionary of agents' observations.
        :type dict_agents_obs: Dict[str, Any]
        :param infos: Additional information.
        :type infos: Dict[str, Dict[str, Any]]
        :param is_last_timestep: Flag indicating if it is the last timestep, defaults to False.
        :type is_last_timestep: bool, optional
        :return: Multi-agent observation, updated infos, and a flag indicating if it is the lowest level.
        :rtype: Tuple[Dict[str, Any], Dict[str, Dict[str, Any]], bool]
        """
        is_lowest_level = True
        return dict_agents_obs, infos, is_lowest_level
    
    def set_low_level_obs(
        self,
        env_config: Dict[str, Any],
        agent_states: Dict[str,Dict[str,Any]],
        dict_agents_obs: Dict[str,Any],
        infos: Dict[str, Dict[str, Any]],
        goals: Dict[str, Any]
    ) -> Tuple[Dict[str, Any], Dict[str, Dict[str, Any]], bool]:
        """
        Set the multiagent observation.

        :param env_config: environment configuration
        :type env_config: Dict[str,Any]
        :param agent_states: agent states
        :type agent_states: Dict[str,Any]
        :param dict_agents_obs: dictionary of agents observations
        :type dict_agents_obs: Dict[str,Any]
        :return: multiagent observation
        :rtype: Dict[str,Any]
        """
        is_lowest_level = True
        return dict_agents_obs, infos, is_lowest_level
