#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/7/18
file:       env.py
version:
doc:        环境规范声明.
"""
import abc
from typing import Any, Optional, Tuple, Dict

from ray.rllib.env.env_context import EnvContext
from ray.rllib.env.multi_agent_env import MultiAgentEnv

from corekit.agents.base import Agent
from corekit.utils.typings import AgentConfig, AgentID, MultiAgentDict


class BaseEnv(MultiAgentEnv, abc.ABC):
    """环境基类"""
    _skip_env_checking: bool = False  # 是否跳过环境规范检查

    count_episode: int = -2  # 回合数记录
    count_step: int = 0  # 步进数记录
    worker_index: int = 1  # 环境的worker序号, 从 1 开始

    agents: dict[AgentID, Agent]
    env_config: dict  # 环境配置
    engine: Any  # 引擎模块
    scenario_path: str  # 引擎想定文件的路径

    def __init__(self, env_config: "EnvContext"):
        super().__init__()
        self.env_config: dict = env_config.copy()
        self.worker_index = env_config.worker_index

        if getattr(self, '_skip_env_checking', False) or env_config.get('disable_env_checking'):
            self.count_episode = -1

    def _init(self, scenario_path: str = None):
        """初始化, 包括处理想定等, 生成agents等功能"""

    def get_agents(self, agent_configs: dict[AgentID, AgentConfig] = None) -> dict[AgentID, Agent]:
        """生成智能体对象

        这个功能之所以放到env里, 是因为智能体的数量和设置应该和引擎想定有关
        """
        agents = {
            aid: Agent(agent_config)
            for aid, agent_config in agent_configs.items()
        }
        return agents

    def reset(
        self,
        *,
        seed: Optional[int] = None,
        options: Optional[dict] = None,
    ) -> Tuple[MultiAgentDict, MultiAgentDict]:
        super().reset(seed=seed, options=options)
        self.count_episode += 1
        self.count_step = 0

        return {}, {}

    def step(
        self, action_dict: MultiAgentDict
    ) -> Tuple[MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict]:
        if not action_dict:  # 没声明 action_space 时会产生空指令字典
            return {}, {}, {'__all__': False}, {'__all__': False}, {}

        # 动作预处理
        action_excecute = {
            aid: self.agents[aid].act_trans(act_array)
            for aid, act_array in action_dict.items()
        }
        self.engine.step(action_excecute)

    def reward(self):
        """MBMPO训练算法需要该函数"""
        reward = {agent_id: 1. for agent_id in self._agent_ids}
        return reward


class BaseHierarchicalEnv(BaseEnv):
    """分层环境基类"""
    _skip_env_checking: bool = True

    all_agents: dict[AgentID, Agent]

    def get_agents(self, agent_configs: dict[AgentID, AgentConfig] = None):  # -> dict[AgentID, Agent]:
        raise NotImplementedError

    def step(
        self, action_dict: MultiAgentDict
    ) -> Tuple[MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict, MultiAgentDict]:
        """同一层级必须同时决策"""
        # if ...:
        #     raise NotImplementedError

        raise NotImplementedError
