
from .base_agent import BaseAgent

import time

from ..utils.chat_template import Query

import json

class ReactAgent(BaseAgent):
    def __init__(self,
                 agent_name,
                 task_input,
                 agent_process_factory,
                 log_mode: str
        ):
        BaseAgent.__init__(
            self,
            agent_name,
            task_input,
            agent_process_factory,
            log_mode
        )

        # self.tool_list = {}

        self.plan_max_fail_times = 3 # 这个变量表示在生成计划时，允许计划失败的最大次数
        self.tool_call_max_fail_times = 3 # 这个变量表示在调用工具时，允许工具调用失败的最大次数

    def build_system_instruction(self):
        prefix = "".join( # prefix 提供了初始的系统指令，描述了 agent 的总体任务或背景信息。
            [
                "".join(self.config["description"])
            ]
        )
        plan_instruction = "".join( # plan_instruction 提供了详细的步骤和格式，指导 agent 如何生成操作计划并在每一步中选择最相关的工具。
            [
                f'You are given the available tools from the tool list: {json.dumps(self.tools)} to help you solve problems.',
                'Generate a plan of steps you need to take.',
                'The plan must follow the json format as: ',
                '[',
                '{"message": "message_value1","tool_use": [tool_name1, tool_name2,...]}',
                '{"message": "message_value2", "tool_use": [tool_name1, tool_name2,...]}',
                '...',
                ']',
                'In each step of the planned workflow, you must select the most related tool to use',
                'Followings are some plan examples:',
                '[',
                '{"message": "Gather information from arxiv", "tool_use": ["arxiv"]},',
                '{"message", "Based on the gathered information, write a summarization", "tool_use": []}',
                '];',
                '[',
                '{"message": "identify the tool that you need to call to obtain information.", "tool_use": ["imdb_top_movies", "imdb_top_series"]},',
                '{"message", "based on the information, give recommendations for the user based on the constrains.", "tool_use": []}',
                '];',
                '[',
                '{"message": "identify the tool that you need to call to obtain information.", "tool_use": ["imdb_top_movies", "imdb_top_series"]},',
                '{"message", "based on the information, give recommendations for the user based on the constrains.", "tool_use": []}',
                '];',
                '[',
                '{"message": "identify the tool that you need to call to obtain information.", "tool_use": ["imdb_top_movies", "imdb_top_series"]},'
                '{"message", "based on the information, give recommendations for the user based on the constrains.", "tool_use": []}',
                ']'
            ]
        )
        # exection_instruction = "".join( # 这个到底有没有执行
        #     [
        #         'To execute each step in the workflow, you need to output as the following json format:',
        #         '{"[Action]": "Your action that is indended to take",',
        #         '"[Observation]": "What will you do? If you will call an external tool, give a valid tool call of the tool name and tool parameters"}'
        #     ]
        # )
        if self.workflow_mode == "manual":
            self.messages.append(
                {"role": "system", "content": prefix}
            )
            # self.messages.append( # 这个到底有没有执行
            #     {"role": "user", "content": exection_instruction}
            # )
        else:
            assert self.workflow_mode == "automatic"
            self.messages.append(
                {"role": "system", "content": prefix}
            )
            self.messages.append(
                {"role": "user", "content": plan_instruction}
            )

    def automatic_workflow(self):
        return super().automatic_workflow()

    def manual_workflow(self):
        pass

    def call_tools(self, tool_calls):
        # self.logger.log(f"***** It starts to call external tools *****\n", level="info")
        success = True
        actions = []
        observations = []
        for tool_call in tool_calls:
            function_name = tool_call["name"]
            function_to_call = self.tool_list[function_name]
            function_params = tool_call["parameters"]

            try:
                function_response = function_to_call.run(function_params)
                actions.append(f"I will call the {function_name} with the params as {function_params}")
                observations.append(f"The knowledge I get from {function_name} is: {function_response}")

            except Exception:
                actions.append("I fail to call any tools.")
                observations.append(f"The tool parameter {function_params} is invalid.")
                success = False

        return actions, observations, success

    """
        构建系统指令并设置初始任务输入。
        确定工作流程模式（自动或手动）并执行相应的工作流程。
        处理工作流程中的每一步，处理工具调用并记录结果。
        返回agent的性能摘要，包括最终结果、轮次和时间信息。
    """
    def run(self):
        self.build_system_instruction() # 构建系统指令并设置总任务。

        task_input = self.task_input

        self.messages.append(
            {"role": "user", "content": task_input}
        )
        self.logger.log(f"{task_input}\n", level="info")

        workflow = None

        if self.workflow_mode == "automatic": # 确定工作流程模式（自动或手动）并执行相应的工作流程。
            workflow = self.automatic_workflow()
        else:
            assert self.workflow_mode == "manual"
            workflow = self.manual_workflow()

        self.messages.append(
            {"role": "assistant", "content": f"[Thinking]: The workflow generated for the problem is {json.dumps(workflow)}"}
        )

        self.logger.log("*********************************\n", level="info")
        self.logger.log(f"Generated {self.workflow_mode} workflow is: {workflow}\n", level="info")
        self.logger.log("*********************************\n", level="info")

        if workflow:
            final_result = "" # final_result 变量用于存储最终的结果。

            for i, step in enumerate(workflow): # 处理工作流程中的每一步，处理工具调用并记录结果。
                message = step["message"] # 每一步固定格式，message和tool_use
                tool_use = step["tool_use"]

                prompt = f"At step {self.rounds + 1}, you need to {message}. "
                self.messages.append({
                    "role": "user",
                    "content": prompt
                })
                if tool_use: # 选择工具
                    selected_tools = self.pre_select_tools(tool_use)

                else:
                    selected_tools = None

                response, start_times, end_times, waiting_times, turnaround_times = self.get_response( # 得到回应
                    query = Query(
                        messages = self.messages,
                        tools = selected_tools
                    )
                )
                if self.rounds == 0:
                    self.set_start_time(start_times[0])

                # execute action
                response_message = response.response_message

                tool_calls = response.tool_calls

                self.request_waiting_times.extend(waiting_times)
                self.request_turnaround_times.extend(turnaround_times)

                if tool_calls:
                    for _ in range(self.plan_max_fail_times):
                        actions, observations, success = self.call_tools(tool_calls=tool_calls) # observation是调用工具以后观察到的东西

                        action_messages = "[Action]: " + ";".join(actions) # "I will call the tool_name with the params as function_params"
                        observation_messages = "[Observation]: " + ";".join(observations) # "The knowledge I get from tool_name is: function_response" or "The tool parameter function_params is invalid."

                        self.messages.append(
                            {
                                "role": "assistant",
                                "content": action_messages + ";;" + observation_messages
                            }
                        )
                        if success:
                            break
                else:
                    thinkings = response_message
                    self.messages.append({
                        "role": "assistant",
                        "content": thinkings
                    })

                if i == len(workflow) - 1:
                    final_result = self.messages[-1]

                self.logger.log(f"At step {i + 1}, {self.messages[-1]}\n", level="info")

                self.rounds += 1

            self.set_status("done")
            self.set_end_time(time=time.time())

            return { # 返回agent的性能摘要，包括最终结果、轮次和时间信息。
                "agent_name": self.agent_name,
                "result": final_result,
                "rounds": self.rounds,
                "agent_waiting_time": self.start_time - self.created_time,
                "agent_turnaround_time": self.end_time - self.created_time,
                "request_waiting_times": self.request_waiting_times,
                "request_turnaround_times": self.request_turnaround_times,
            }

        else:
            return {
                "agent_name": self.agent_name,
                "result": "Failed to generate a valid workflow in the given times.",
                "rounds": self.rounds,
                "agent_waiting_time": None,
                "agent_turnaround_time": None,
                "request_waiting_times": self.request_waiting_times,
                "request_turnaround_times": self.request_turnaround_times,
            }
