from enum import Enum
from itertools import chain
from typing import Annotated

from judy_core.agents.prompts import AGENT_COMMON_INSTRUCTION, AGENT_FOLDER_INSTRUCTION
from judy_core.agents.qa_agent import QAAgent
from judy_core.agents.react_agent import ReactAgent
from judy_core.agents.tools import ASK_MEMBER_TOOL_NAME, AskMemberTool, AskTargets
from judy_core.events import CreateAgentFolder
from judy_core.messages import (
    ExecutionEndReason,
    ExecutorResult,
    Subtask,
    TaskHistory,
    TeamMemberId,
)
from judy_core.utils.tool import TOOL_RESULT_SUCCESS_CONTENT, Tool, ToolRegistry
from pydantic import Field
from tongcell.core import ActorId, MessageContext, message_handler
from tongcell.util import getLogger

END_EXECUTION_TOOL_NAME = "end_execution"

EXECUTOR_NOT_ISSUE_ANY_TOOL_CALL_INSTRUCTION = f"""如果你认为当前已经可以进入汇报阶段，你需要调用{END_EXECUTION_TOOL_NAME}工具；否则你应该调用其他工具继续执行。"""

EXECUTOR_FEEDBACK_INSTRUCTION = "你需要就目前的发现和取得的进展进行汇报，汇报中可以包含你的建议。该汇报会被提交给planner，以便其及时调整规划。"

EXECUTOR_FEEDBACK_PHASE_INSTRUCTION = f"""\
你现在进入了汇报阶段，{EXECUTOR_FEEDBACK_INSTRUCTION}**你不能再调用任何工具了，直接将汇报作为你的回复内容即可**。\
"""

EXECUTOR_MAX_STEPS = 80
EXECUTOR_STEPS_EXAHUSTED_INSTRUCTION = f"你已经执行到了第{EXECUTOR_MAX_STEPS}步，超过了最大步数限制，你不能再继续了。{EXECUTOR_FEEDBACK_PHASE_INSTRUCTION}"


TASK_HISTORY_INSTRUCTION = f"""\
# 任务历史
任务历史由团队成员的执行结果组成，按照成员的执行顺序排列。

executor的执行结果主要包含其负责的todo，汇报内容等:
- step_count表示其执行步数，tool_call_count表示其工具调用次数，因为一步可以调用多个工具，所以tool_call_acount可能大于step_count。
- 部分executor会一次性完成多个todo。
- 由于planner会根据反馈调整规划，所以executor执行过的一些todo并不包含在当前规划中。

verifier的执行结果主要包含其验证报告。

## 善用提问
`{ASK_MEMBER_TOOL_NAME}` tool可用于向其他团队成员（任务历史中包含有团队成员的名字）提问。

提问的用途包括但不限于：
- 解决任务历史中团队成员的报告内容不够清晰，过于简略，或者由于内容过长被截断的问题
- 收集团队成员的成功/失败经验，避免重复工作，提高效率
- 在verifier给出验收失败的结论后，向其咨询验收的方式

注意：
- 执行步数/工具调用次数(step count > 10)越多的成员，向其提问价值越高，因为其汇报内容可能并没有完全覆盖他发现的所有信息和尝试过的所有路径。
- 团队成员都非常热心，会诚实和毫无保留地回答你的问题。
- 你可以通过一次性向多个人提问来提高效率。
"""


EXECUTOR_INSTRUCTION = f"""\
你的团队通过协作来尝试完成一项编程任务，团队由1个planner，1个verifier，和一组executor构成。你是其中一个executor。

# 对你的要求

## 分工
planner负责制定规划（由todo列表构成），为当前规划中第一个处于PENDING状态的todo分配executor，executor使用为其提供的工具来解决这个todo。
executor配备有相同的tool，像接力赛一样接力处理规划中的todo。你是这些executor其中之一。

## **你可以向planner提建议**
- 由于planner不知道你有哪些工具，以及你的工具的适用范围。可能也不知道作为一线人员的executor所发现的细节，所以planner给出的plan或者todo不一定是最优的。
- 如果根据你的系统指令、工具的说明、以及你所发现的信息，你认为当前规划/todo并不合理，或者你发现有更优/更合适的方案时，你可以中止执行，直接进行任务汇报，告诉planner你的想法和建议。
- 你的建议可以不局限于你当前所负责的todo，也可以包含已经执行过的todo或者整体规划的建议。

{AGENT_FOLDER_INSTRUCTION}

# 工具

## 并行工具调用
你可以在一次回复中包含多个工具调用（tool calls），这些工具调用会被依次执行。你必须充分利用此特性来减少交互轮数，降低任务耗时。
此外如非必要，不要使用shell的`&&`或`;`或`\n`语法来在一个工具调用中串联多个命令，而是应该拆分为多个工具调用，在一次回复中一次性返回。这样多个命令的输出不会混在一起，对你理解工具的执行结果更有帮助。

## 尽量通过并行来减少步数，减少任务耗时
除了并行工具调用外，你还可以通过其它方式（比如批量运行）来提高效率。比如你在编写脚本时，不要一个一个地去验证或者测试，而是编写一个包含多条数据或者用例的脚本，批量验证。
**在进行parallel tool call时，你必须确保生成调用llm的input文件的tool call，位于调用llm 的tool call之前。**

## 工具调用的超时时间
在运行命令或者等待命令退出时，超时时间最好设置在120秒以下，定期观察进度，有助于及时发现和处理死锁、卡顿等异常问题。
你可以在等待期间尝试其它方案，思考进程的耗时是否符合预期。

**如果代码运行很久都没有出结果，有可能是因为代码存在bug**，如：
- 死循环、无限递归、时间复杂度过高等计算量过大的问题
- 死锁类问题
你可以通过查看/proc目录下进程的cpu使用时间的变化来确定这类问题。

## 工具的输出可能会被截断
如果工具的输出内容过长，为了避免读取到大量垃圾内容，会被截断：
- 行内截断：当某行内容过长时，行中间部分会被截断，以 `...[TRUNCATED yyy chars]...` 来表示。
- 中间行截断：经过行内截断处理后，如果整体内容仍然过长，则内容头部、尾部各保留一部分，中间部分的行会被整体截断，以 `...[TRUNCATED xxx lines, yyy chars]...` 表示

### 截断处理策略
如果被截断的内容可能含有关键信息，那么你需要进行过滤，在减少无意义内容的同时保留关键信息。 比如：
- 对于行内截断：使用`grep -o`来仅输出匹配内容
- 对于中间行截断：使用`grep`,`grep-v`来保留和过滤整行内容。
如果简单的过滤命令无法满足需求，你可以编写脚本来实现复杂的过滤逻辑 。

**过滤容易导致意外错过重要信息**
- 过滤策略应该从宽松逐渐滑向严格，不要直接转向最严格的策略
- 不要只使用一种搜索/过滤策略，最好正向、反向、宽松、严格等多种策略都尝试一下
为了避免遗漏信息，你应该综合交叉验证后再下结论。

## 调试(debug)
排查bug（或者coredump）时，相比盲猜和不断地改代码加日志，通过debugger（如gdb、pdb）定位问题更加高效和准确。
可以通过在interactive shell中启动debugger来交互式地调试代码。

## 网页搜索
如果你拥有搜索网页的tool，或者拥有可以开启网页搜索功能调用外部llm的tool，你应该优先使用这些tool来搜索网页。
这些tool的搜索质量和效率一般情况下比你自行写代码爬取搜索页面要好很多，且没有限流、反爬虫等问题。

### 把搜索tool当成你的百科全书
**如果你对某些领域/知识/信息不是很了解，或者拿不准，不要瞎猜，务必通过搜索来补充相关知识。**

## 多模态
对于音视频、图片或者布局复杂/非结构化的文档（如pdf），相比于简单的OCR，多模态能够额外提供
- 文档布局
- 形状、图案、颜色、空间关系
- 扭曲文字的识别
等信息。

多模态能力的用途包括但不限于：
- 分析图片，视频，pdf中的特征：通过分析采样数据（如抽取视频的某一帧，从文件夹中选择一个pdf/图片文件），形成数据的大致轮廓，精准制定规划/确定下一步的方向。
- 将其他格式的文件渲染为图片，进行可视化分析
- debug/检查/验证/查漏补缺：对使用常规手段的批处理结果中的可疑数据或者抽样数据，使用多模态能力进行二次检验。

如果你能够调用外部LLM的话，你还可以使用外部LLM的多模态能力。

# 自身和外部LLM的文本/多模态能力的使用原则

## 禁用自身和外部LLM的文本/多模态能力的情形
### 需要准确的颜色/文本等信息时
你自身的多模态能力，以及外部LLM的文本/多模态能力，有时可能存在幻觉，并不准确，仅供参考。
当需要非常准确的信息时，应该使用专用的命令/SDK/模型（当某个命令/SDK模型不可用时，不要放弃，可以尝试以下同类型的其它替代品)。
### 数据量过大时
当数据量过大（如文本大于500k）或者文件过多（大于30个）时，禁止使用自身或者外部LLM的文本/多模态能力来处理数据，因为**时间、金钱成本会非常高昂**。
这时你必须使用其它文本或音视频处理手段（如OCR等)，但可以使用多模态能力来对采样结果进行分析，或者对部分处理结果进行验证。

## 其它情况下应优先使用自身多模态能力（相比于调用外部LLM）：
仅在以下情形考虑调用外部LLM的多模态能力：
- 你自身没有多模态能力
- 需要处理的数据量过大（如文本大于100k）或者文件过多（大于10个）
- 需要外部LLM的结果作为参考或者交叉验证
因为相比于调用外部LLM，使用你自身的多模态能力有以下优势：
- 保留更多细节信息，避免外部LLM转述造成信息丢失
- 减少2轮交互，节省时间和金钱成本，提高效率

# 效率
任务的处理效率非常重要，因此你必须充分使用缓存、批处理等技巧来提高效率。

## 高耗时环节的处理原则
当运行的某个命令，或者代码中的某个步骤，由于计算量较大或者需要访问网络等原因，耗时可能会很高时，必须遵循以下原则：

### 及时保存中间处理结果
为了避免重复执行耗时环节，应及时保存中间处理结果。典型场景包括：
- 批处理中的部分数据处理出错，导致整个批次需要重跑
- 后续处理环节执行失败，需要重新运行前面的环节
- 需要利用中间结果进行验证、调试或分析
通过保存中间结果，可以避免重复计算，提高效率。

### 周期性输出处理进度
对于耗时较长的环节，应定期输出处理进度（最好带上时间戳），以便
- 了解当前进展
- 预估整体任务的剩余执行时间
- 及时发现卡顿、死锁等异常情况
同时注意，处理进度的输出不应该过于频繁，以免产生过多无用信息，导致关键信息被淹没。
可以选择间隔一定时间（如每隔10秒）或者处理完一定量的数据（如每处理30个文件）后输出一次进度。

### 输出/打印处理进度、日志、数据等关键信息时，务必及时flush
如果输出时不进行flush操作，输出内容可能会一直在进程内部缓存，直到进程正常退出后才能获取到。
对于运行时间较长的命令/代码，这会导致无法及时看到相关输出，严重影响开发、诊断和调试效率。

## 必要时可以手动处理corner case
一些常规方法，如：
- 基于关键字或者正则表达式的文本处理方法
- 基于小模型的文本/图片/视频处理方法
经常会遇到难以处理的corner case，当corner case数据量较小时，使用多模态能力手动处理/修正，相比不断优化代码、训练模型花费的成本（时间、金钱）更低。


# 你的几个阶段
## 执行阶段
- 这是你的初始阶段，在这个阶段你使用工具来尝试执行todo。
- 当你认为执行阶段可以结束时（比如当前todo完成了，或者已经无法取得进展了，或者你想建议planner修改规划和todo时）时，你需要调用{END_EXECUTION_TOOL_NAME}工具。
- 该工具被调用后然后会进入汇报阶段。 
- **你只能在执行阶段调用工具，在之后的阶段（如汇报和QA）里你不能调用任何工具。**
每个executor最多可以执行{EXECUTOR_MAX_STEPS}步，超过限制后系统会剥夺其工具执行权，立即进入汇报阶段。

## 汇报阶段
在汇报阶段，{EXECUTOR_FEEDBACK_INSTRUCTION}。

## QA阶段
汇报阶段结束后，你会进入QA阶段，你会收到来自其他人的问题。

{TASK_HISTORY_INSTRUCTION}

{AGENT_COMMON_INSTRUCTION}
"""

# 注意：**此阶段可以和执行阶段合并在一起，即你可以在调用{END_EXECUTION_TOOL_NAME}工具的同时，\
# 直接将你的汇报内容作为常规的文本消息返回(注意不是放在tool call里，而是和tool_call并列），这样可以减少一次交互，提高效率。**

logger = getLogger(__name__)


class ExecutorState(Enum):
    READY = 0
    WAITING_TOOL_RESULT = 1


def get_ask_targets(task_history: TaskHistory) -> AskTargets:
    ask_targets: AskTargets = {}
    for member_result in task_history.team_member_results:
        member_id = member_result.member_id
        ask_targets[member_id.name] = member_id

    return ask_targets


class Executor(QAAgent, ReactAgent):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self._state = ExecutorState.READY

        self._execution_end_reason: ExecutionEndReason | None = None
        self._extra_finished_todo_indexes: list[int] = []

    def _react_exit_prompt(self) -> str:
        return EXECUTOR_NOT_ISSUE_ANY_TOOL_CALL_INSTRUCTION

    @message_handler()
    async def start_execute(self, subtask: Subtask, ctx: MessageContext) -> None:
        self._sub_task = subtask
        self._step_count = 0
        self._tool_call_count = 0

        ask_member_tool = AskMemberTool(ctx)
        ask_member_tool.set_ask_targets(get_ask_targets(self._sub_task.task_history))
        self._internal_tool_registry = ToolRegistry(
            [
                Tool.from_callable(ask_member_tool, ASK_MEMBER_TOOL_NAME),
                Tool.from_callable(self.end_execution, END_EXECUTION_TOOL_NAME),
            ],
        )
        self._tools = [
            tool_def.to_openai_format()
            for tool_def in chain(
                self._model_config.tools_definition,
                self._internal_tool_registry.get_tools_definition(),
            )
        ]
        # if self._need_explicit_cache() and self._tools:
        #     self._tools[-1]["cache_control"] = {"type": "ephemeral"}

        personal_dir = self._id.name
        await self._publish_output_event(CreateAgentFolder(folder_name=personal_dir))

        prompt = "\n".join(
            [
                subtask.task.to_prompt(),
                self._task_config.to_prompt(personal_dir),
                "为了避免多段文本和格式造成混淆，部分内容用xml格式表示。",
                "# 任务执行历史",
                subtask.task_history.to_prompt(),
                "",
                "# planner的当前规划",
                subtask.plan.to_prompt(),
                "",
                f"你的名字是`{self._id.name}`，你现在负责规划中的第{subtask.todo_index + 1}项，即：{subtask.plan.todos[subtask.todo_index].desc}",
            ]
        )

        executor_instruction = EXECUTOR_INSTRUCTION
        if subtask.task.env_notes:
            executor_instruction += subtask.task.env_notes

        system_part: dict = {
            "type": "text",
            "text": executor_instruction,
        }
        if self._need_explicit_cache():
            system_part["cache_control"] = {"type": "ephemeral"}
        self._messages = [
            {
                "role": "system",
                "content": [system_part],
            },  # pyright: ignore[reportAttributeAccessIssue]
            {"role": "user", "content": prompt},
        ]

        await self._publish_model_input(self._messages)
        await self._do_turn(ctx)

    async def _on_turn_done(self, ctx: MessageContext) -> None:
        # TODO：限制最大token数
        if self._execution_end_reason is None:
            self._step_count += 1
            if self._step_count < EXECUTOR_MAX_STEPS:
                await self._do_turn(ctx)
            else:
                self._execution_end_reason = ExecutionEndReason.STEPS_EXHAUSTED
                report_content = await self._gen_report(True, ctx)
                await self._send_report(report_content, ctx)
        else:
            # if self._tool_message.content:
            #     report_content = self._tool_message.content
            # else:
            report_content = await self._gen_report(False, ctx)
            await self._send_report(report_content, ctx)

    async def _gen_report(self, steps_exhausted: bool, ctx: MessageContext) -> str:
        if steps_exhausted:
            feed_back_instruction = EXECUTOR_STEPS_EXAHUSTED_INSTRUCTION
            self._messages.append(
                {
                    "role": "user",
                    "content": feed_back_instruction,
                }
            )
            cache_control_injection_points = None
        else:
            cache_control_injection_points = [
                {
                    "location": "message",
                    "index": -3,
                }
            ]

        choice = await self._acompletion_without_tool_calls(
            messages=self._messages,
            tools=self._tools,
            cache_control_injection_points=cache_control_injection_points,
        )
        self._messages.append(choice.message)
        assert choice.message.content
        return choice.message.content

    async def _send_report(self, report_content: str, ctx: MessageContext):
        from judy_core.agents.planner import Planner

        assert self._execution_end_reason is not None

        await ctx.send(
            ExecutorResult(
                member_id=TeamMemberId.from_actor_id(self._id),
                end_reason=self._execution_end_reason,
                plan=self._sub_task.plan,
                todo_index=self._sub_task.todo_index,
                extra_finished_todo_indexes=self._extra_finished_todo_indexes,
                report_content=report_content,
                step_count=self._step_count,
                tool_call_count=self._tool_call_count,
            ),
            ActorId.default(Planner.actor_kind),
        )

    async def end_execution(
        self,
        reason: ExecutionEndReason,
        extra_finished_todo_indexes: Annotated[
            list[int] | None,
            Field(
                description="如果你本次执行期间不仅完成了当前todo，还完成了其他todo，可以把他们的序号加到这里"
            ),
        ] = None,
    ):
        """当你认为你可以停止执行的时候（即当你负责的todo已经完成，或者失败，或者你想向planner提出建议时），调用这个工具。

        该工具被调用后会立刻进入汇报阶段，因此你只能在最后一步才能调用这个工具。否则你将被提前中止执行，立即进入汇报阶段。
        """
        self._execution_end_reason = reason

        self._extra_finished_todo_indexes = []
        if (
            reason == ExecutionEndReason.COMPLETED
            and extra_finished_todo_indexes is not None
        ):
            for extra_finished_todo_index in extra_finished_todo_indexes:
                extra_finished_todo_index -= 1
                if extra_finished_todo_index == self._sub_task.todo_index:
                    continue
                if extra_finished_todo_index < 0:
                    continue
                if extra_finished_todo_index >= len(self._sub_task.plan.todos):
                    continue
                self._extra_finished_todo_indexes.append(extra_finished_todo_index)

        return TOOL_RESULT_SUCCESS_CONTENT + ". " + EXECUTOR_FEEDBACK_PHASE_INSTRUCTION
