if TYPE_CHECKING:
    from agentverse.agents.base import BaseAgent
    from agentverse.message import Message


class BaseRule(BaseModel):
    pass


class BaseEnvironment(BaseModel):
    """
    Base class for environment.

    Args:
        agents: List of agents
        rule: Rule for the environment
        max_turns: Maximum number of turns
        cnt_turn: Current turn number
        last_messages: Messages from last turn
        rule_params: Variables set by the rule
    """

    agents: List[BaseAgent]
    rule: BaseRule
    max_turns: int = 10
    cnt_turn: int = 0
    last_messages: List[Message] = []
    rule_params: Dict = {}

    @abstractmethod
    async def step(self) -> List[Message]:
        """Run one step of the environment"""
        pass

    @abstractmethod
    def reset(self) -> None:
        """Reset the environment"""
        pass

    def report_metrics(self) -> None:
        """Report useful metrics"""
        total_spent = sum([agent.get_spend() for agent in self.agents])
        logger.info(f"Total spent: ${total_spent}")
        pass

    def is_done(self) -> bool:
        """Check if the environment is done"""
        return self.cnt_turn >= self.max_turns

class BasicEnvironment(BaseEnvironment):

    # 导入Tasksolving_rule 
    # Tasksolving内定义了如下规则：role_assigner（角色分配）decision_maker（决策制订）executor（执行）evaluator（评估）
    rule: TasksolvingRule
    agents: Dict[Enum, Union[BaseAgent, List[BaseAgent]]] = None

    task_description: str

    cnt_turn: int = 0
    max_turn: int = 10
    success: bool = False

    # 初始化环境
    def __init__(self, **kwargs):
        rule_config = kwargs.pop("rule", {})
        role_assigner_config = rule_config.pop(
            "role_assigner", {"type": "role_description"}
        )
        decision_maker_config = rule_config.pop("decision_maker", {"type": "vertical"})
        executor_config = rule_config.pop("executor", {"type": "none"})
        evaluator_config = rule_config.pop("evaluator", {"type": "basic"})
        rule = TasksolvingRule(
            role_assigner_config=role_assigner_config,
            decision_maker_config=decision_maker_config,
            executor_config=executor_config,
            evaluator_config=evaluator_config,
        )
        super().__init__(rule=rule, **kwargs)
    # step 方法用于执行
    async def step(
        self, advice: str = "No advice yet.", previous_plan: str = "No solution yet."
    ) -> List[Message]:
        result = ""
        logs = []
        # 记录当前回合
        logger.info(f"Loop Round {self.cnt_turn}")

        # ================== EXPERT RECRUITMENT ==================
        # 分配角色，根据（任务描述）task_description （智能体）agents （目前回合）cnt_turn （建议）advice 分配
        agents = await self.rule.role_assign(
            self.task_description, self.agents, self.cnt_turn, advice
        )
        description = "\n".join([agent.role_description for agent in agents])
        logs.append({"module": "Role Assigner", "content": description})
        logger.info("", f"Role Assignment:\n{description}", Fore.CYAN)
        # ================== EXPERT RECRUITMENT ==================

        # ================== DECISION MAKING ==================
        # 调用decision_maker 中的decision_making方法返回生成的计划
        plan: List[SolverMessage] = await self.rule.decision_making(
            self.task_description, self.agents, previous_plan, advice
        )
        flatten_plan = "\n".join([p.content for p in plan])
        logs.append({"module": "Decision Maker", "content": flatten_plan})
        logger.info("", f"Decision Plan:\n{flatten_plan}", Fore.YELLOW)
        # ================== DECISION MAKING ==================

        # ================== EXECUTION ==================
        # 按照plan执行一轮动作
        result: List[ExecutorMessage] = await self.rule.execute(
            self.task_description, self.agents, plan
        )
        flatten_result = "\n".join([r.content for r in result])
        logs.append({"module": "Executor", "content": flatten_result})
        logger.info("", f"Execution Result:", Fore.GREEN)
        logger.info("", flatten_result, Fore.GREEN)
        # ================== EXECUTION ==================

        # ================== EVALUATION ==================
        # 执行结束后由evaluate来生成对任务的评价与后续优化的建议
        score, advice = await self.rule.evaluate(
            self.task_description, self.agents, plan, result
        )
        logs.append(
            {
                "agent": "evaluator",
                "content": f"Evaluation result: Score: {score}\nAdvice: {advice}",
            }
        )
        logger.info(
            "", f"Evaluation result:\nScore: {score}\nAdvice: {advice}", Fore.YELLOW
        )

        if score is not None and (
            # 单任务结果为完成
            (isinstance(score, bool) and score is True)
            # 多任务打分均在 8 以上 为完成
            or (isinstance(score, (list, tuple)) and all([s >= 8 for s in score]))
        ):
            # TODO: 8 is an arbitrary threshold
            logs.append({"agent": "system", "content": "Good score! Accept!"})
            logger.info(
                "", f"Good score! Accept! Final Result:\n{flatten_plan}", Fore.GREEN
            )
            self.success = True
        else:
            logs.append({"agent": "system", "content": "Bad score! Reject!"})
            logger.info("", "Bad score! Reject!", Fore.RED)
        self.cnt_turn += 1
        return flatten_result, advice, flatten_plan, logs, self.success
    # 遍历 agents
    def iter_agents(self):
        for role, agent_or_agents in self.agents.items():
            if isinstance(agent_or_agents, list):
                for agent in agent_or_agents:
                    yield role, agent
            else:
                yield role, agent_or_agents
    # 获取每个agent 的token花费
    def get_spend(self):
        total_spent = sum([agent.get_spend() for (_, agent) in self.iter_agents()])
        return total_spent
    # report token 花费
    def report_metrics(self) -> None:
        logger.info("", "Agent spend:", Fore.GREEN)
        for role, agent in self.iter_agents():
            name = agent.name.split(":")[0]
            logger.info(
                "",
                f"Agent (Role: {role}) {name}: {agent.get_spend_formatted()}",
                Fore.GREEN,
            )
        logger.info("", f"Total spent: ${self.get_spend():.6f}", Fore.GREEN)

    def is_done(self):
        """Check if the environment is done"""
        return self.cnt_turn >= self.max_turn or self.success

    def set_task_description(self, task_description: str = ""):
        self.task_description = task_description

    # 重置环境
    def reset(self) -> None:
        """Reset the environment"""
        self.cnt_turn = 0
        self.rule.reset()

class TasksolvingRule(BaseRule):
    # role_assigner（角色分配）decision_maker（决策制订）executor（执行）evaluator（评估）
    role_assigner: BaseRoleAssigner
    decision_maker: BaseDecisionMaker
    executor: BaseExecutor
    evaluator: BaseEvaluator


    role_assign_only_once: bool = False
    add_execution_result_to_critic: bool = False
    add_execution_result_to_solver: bool = False

    def __init__(
        self,
        role_assigner_config,
        decision_maker_config,
        executor_config,
        evaluator_config,
        *args,
        **kwargs,
    ):
        # 根据yaml 配置文件内内容初始化组件
        def build_components(config: Dict, registry):
            component_type = config.pop("type")
            component = registry.build(component_type, **config)
            return component

        role_assigner = build_components(
            role_assigner_config,
            role_assigner_registry,
        )
        decision_maker = build_components(
            decision_maker_config,
            decision_maker_registry,
        )
        executor = build_components(executor_config, executor_registry)
        evaluator = build_components(evaluator_config, evaluator_registry)
        super().__init__(
            role_assigner=role_assigner,
            decision_maker=decision_maker,
            executor=executor,
            evaluator=evaluator,
            *args,
            **kwargs,
        )

    # 角色分配，初始化agent
    async def role_assign(
        self,
        task_description: str,
        agents: List[BaseAgent],
        cnt_turn: int,
        advice: str = "",
    ) -> List[BaseAgent]:
        """Assign roles to agents"""
        if self.role_assign_only_once and cnt_turn > 0:
            agents = [agents[AGENT_TYPES.SOLVER]] + agents[AGENT_TYPES.CRITIC]
        else:
            agents = await self.role_assigner.astep(
                role_assigner=agents[AGENT_TYPES.ROLE_ASSIGNMENT],
                group_members=[agents[AGENT_TYPES.SOLVER]] + agents[AGENT_TYPES.CRITIC],
                advice=advice,
                task_description=task_description,
            )
            if self.role_assign_only_once and cnt_turn == 0:
                agents[AGENT_TYPES.SOLVER] = agents[0]
                agents[AGENT_TYPES.CRITIC] = agents[1:]
        return agents

    # 生成决策
    async def decision_making(
        self,
        task_description: str,
        agents: List[BaseAgent],
        previous_plan: str,
        advice: str = "No advice yet.",
    ) -> List[SolverMessage]:
        # TODO: plan should be string or a special type of object?

        # dynamic
        if "dynamic" in self.decision_maker.name:
            plan = await self.decision_maker.astep(
                agents=[agents[AGENT_TYPES.SOLVER], *agents[AGENT_TYPES.CRITIC]],
                manager=agents[AGENT_TYPES.MANAGER],
                task_description=task_description,
                previous_plan=previous_plan,
                advice=advice,
            )
        else:
            plan = await self.decision_maker.astep(
                agents=[agents[AGENT_TYPES.SOLVER], *agents[AGENT_TYPES.CRITIC]],
                task_description=task_description,
                previous_plan=previous_plan,
                advice=advice,
            )
        return plan

    # 执行
    async def execute(
        self,
        task_description: str,
        agents: List[BaseAgent],
        final_solution: List[SolverMessage],
    ) -> Any:
        """execution stage.
        Use the executor to finish the task.
        """

        results = await self.executor.astep(
            agents[AGENT_TYPES.EXECUTION], task_description, final_solution
        )
        if self.add_execution_result_to_critic:
            for agent in agents[AGENT_TYPES.CRITIC]:
                agent.add_message_to_memory(results)
        if self.add_execution_result_to_solver:
            agents[AGENT_TYPES.SOLVER].add_message_to_memory(results)
        return results

    # 评估结果
    async def evaluate(
        self,
        task_description: str,
        agents: List[BaseAgent],
        solution: List[SolverMessage],
        result: List[ExecutorMessage],
    ) -> Tuple[List[int], str]:
        """evaluation stage."""
        # if self.human_eval:
        #     print("This round, LLM gave the following result:")
        #     print(result)
        #     comprehensiveness = input("Please evaluate the comprehensiveness>> ")
        #     detailedness = input("Please evaluate the detailedness>> ")
        #     feasibility = input("Please evaluate the feasibility>> ")
        #     novelty = input("Please evaluate the novelty>> ")
        #     advice = input("Please give some advice>>")
        #     try:
        #         comprehensiveness = int(comprehensiveness)
        #         detailedness = int(detailedness)
        #         feasibility = int(feasibility)
        #         novelty = int(novelty)
        #     except ValueError:
        #         logger.error("Bad response from human evaluator!")
        #     return ([comprehensiveness, detailedness, feasibility, novelty], advice)
        # else:
        evaluation = await self.evaluator.astep(
            agent=agents[AGENT_TYPES.EVALUATION],
            solution=solution,
            result=result,
            task_description=task_description,
            all_role_description=[
                agents[AGENT_TYPES.SOLVER].role_description,
                *[agent.role_description for agent in agents[AGENT_TYPES.CRITIC]],
            ],
        )
        # 返回任务得分，返回评估后的修改建议
        return evaluation.score, evaluation.advice

    def reset(self) -> None:
        self.role_assigner.reset()
        self.decision_maker.reset()
        self.executor.reset()
        self.evaluator.reset()

@EnvironmentRegistry.register("sim-basic")
class BasicEnvironment(BaseEnvironment):
    """
    A basic environment implementing the logic of conversation.

    Args:
        agents: List of agents
        rule: Rule for the environment
        max_turns: Maximum number of turns
        cnt_turn: Current turn number
        last_messages: Messages from last turn
        rule_params: Variables set by the rule
    """

    agents: List[BaseAgent]
    rule: Rule
    max_turns: int = 10
    cnt_turn: int = 0
    last_messages: List[Message] = []
    rule_params: Dict = {}
    # 初始化环境
    def __init__(self, rule, **kwargs):
        rule_config = rule
        order_config = rule_config.get("order", {"type": "sequential"})
        visibility_config = rule_config.get("visibility", {"type": "all"})
        selector_config = rule_config.get("selector", {"type": "basic"})
        updater_config = rule_config.get("updater", {"type": "basic"})
        describer_config = rule_config.get("describer", {"type": "basic"})
        rule = Rule(
            order_config,
            visibility_config,
            selector_config,
            updater_config,
            describer_config,
        )
        super().__init__(rule=rule, **kwargs)
    # 按照规定的agent 顺序执行agent
    async def step(self) -> List[Message]:
        """Run one step of the environment"""

        # Get the next agent index
        agent_ids = self.rule.get_next_agent_idx(self)

        # Generate current environment description
        # 根据不同的agent 生成对应的 env description
        env_descriptions = self.rule.get_env_description(self)
        
        # Generate the next message
        # agent 根据自己对应的描述采取行动
        messages = await asyncio.gather(
            *[self.agents[i].astep(env_descriptions[i]) for i in agent_ids]
        )

        # Some rules will select certain messages from all the messages
        selected_messages = self.rule.select_message(self, messages)
        self.last_messages = selected_messages
        self.print_messages(selected_messages)

        # Update the memory of the agents
        self.rule.update_memory(self)

        # Update the set of visible agents for each agent
        self.rule.update_visible_agents(self)

        self.cnt_turn += 1

        return selected_messages

    def print_messages(self, messages: List[Message]) -> None:
        for message in messages:
            if message is not None:
                # logging.info(f"{message.sender}: {message.content}")
                logger.info(f"{message.sender}: {message.content}")

    def reset(self) -> None:
        """Reset the environment"""
        self.cnt_turn = 0
        self.rule.reset()
        for agent in self.agents:
            agent.reset()

    def is_done(self) -> bool:
        """Check if the environment is done"""
        return self.cnt_turn >= self.max_turns

class SimulationRule(BaseRule):
    """
    Rule for the environment. It controls the speaking order of the agents
    and maintain the set of visible agents for each agent.
    """

    order: BaseOrder
    visibility: BaseVisibility
    selector: BaseSelector
    updater: BaseUpdater
    describer: BaseDescriber

    def __init__(
        self,
        order_config,
        visibility_config,
        selector_config,
        updater_config,
        describer_config,
    ):
        order = order_registry.build(**order_config)
        visibility = visibility_registry.build(**visibility_config)
        selector = selector_registry.build(**selector_config)
        updater = updater_registry.build(**updater_config)
        describer = describer_registry.build(**describer_config)
        super().__init__(
            order=order,
            visibility=visibility,
            selector=selector,
            updater=updater,
            describer=describer,
        )
    # 根据order 获取agent id，按顺序执行
    '''
    class BaseOrder(BaseModel):
    @abstractmethod
    def get_next_agent_idx(self, environment: BaseEnvironment) -> List[int]:
        """Return the index of the next agent to speak"""

    def reset(self) -> None:
        pass
    '''
    def get_next_agent_idx(
        self, environment: BaseEnvironment, *args, **kwargs
    ) -> List[int]:
        """Return the index of the next agent to speak"""
        return self.order.get_next_agent_idx(environment, *args, **kwargs)
    # 
    '''
    class BaseVisibility(BaseModel):
    @abstractmethod
    def update_visible_agents(self, environment: BaseEnvironment):
        """Update the set of visible agents for the agent"""

    def reset(self):
        pass

    '''
    def update_visible_agents(
        self, environment: BaseEnvironment, *args, **kwargs
    ) -> None:
        """Update the set of visible agents for the agent"""
        self.visibility.update_visible_agents(environment, *args, **kwargs)

    # 利用selecor 过滤agent回复的无效信息
    '''
    class BaseSelector(BaseModel):
    @abstractmethod
    def select_message(
        self, environment: BaseEnvironment, messages: List[Message]
    ) -> List[Message]:
        """Selects a set of valid messages from all messages"""
        pass

    def reset(self) -> None:
        pass
    '''
    def select_message(
        self, environment: BaseEnvironment, messages: List[Message], *args, **kwargs
    ) -> List[Message]:
        """Select a set of valid messages from all the generated messages"""
        return self.selector.select_message(environment, messages, *args, **kwargs)

    # 根据事件发生的结果更新agent的记忆（根据visible属性，也即每个agent获取信息的权限）
    '''
    class BaseUpdater(BaseModel):
    """
    The base class of updater class.
    """

    @abstractmethod
    def update_memory(self, environment: BaseEnvironment):
        pass

    def reset(self):
        pass

    '''
    def update_memory(self, environment: BaseEnvironment, *args, **kwargs) -> None:
        """For each message, add it to the memory of the agent who is able to see that message"""
        self.updater.update_memory(environment, *args, **kwargs)

    # 生成对环境的描述
    '''
    class BaseDescriber(BaseModel):
    @abstractmethod
    def get_env_description(
        self, environment: BaseEnvironment, *args, **kwargs
    ) -> List[str]:
        """Return the environment description for each agent"""
        pass

    def reset(self) -> None:
        pass
    '''
    def get_env_description(
        self, environment: BaseEnvironment, *args, **kwargs
    ) -> List[str]:
        """Return the description of the environment for each agent"""
        return self.describer.get_env_description(environment, *args, **kwargs)

    def reset(self) -> None:
        self.order.reset()
        self.visibility.reset()
        self.selector.reset()
        self.updater.reset()
        self.describer.reset()