import time
from copy import copy
from typing import Annotated

from judy_core.agents.ids import VERIFIER_NAME, get_executor_name
from judy_core.agents.llm_agent import LLMAgent
from judy_core.agents.prompts import AGENT_COMMON_INSTRUCTION, AGENT_FOLDER_INSTRUCTION
from judy_core.agents.types import ModelInputMessage
from judy_core.agents.verifier import Verifier
from judy_core.events import DebugTaskLog, Task, TaskFailResult, TaskSuccessResult
from judy_core.messages import (
    ExecutionEndReason,
    ExecutorResult,
    Plan,
    Subtask,
    TaskHistory,
    Todo,
    TodoStatus,
    VerificationRequest,
    VerificationResult,
    VerificationResultStatus,
)
from judy_core.utils.messages import convert_content_to_openai_format
from judy_core.utils.tool import Tool, ToolRegistry
from litellm import BaseModel, Choices
from litellm.types.completion import ChatCompletionToolMessageParam
from pydantic import Field
from tongcell.core import ActorId, MessageContext, message_handler
from tongcell.util import getLogger

SAVE_PLAN_TOOL_NAME = "save_plan"
PLANNER_INSTRUCTION = f"""你是一个编程任务的planner，你与团队中的一组executor协作来尝试完成一项任务。

# 你的角色和职责
你是团队中的架构师（CTO），负责制定规划（plan）。规划由一系列todo构成，你需要保证每个todo清晰、可行，并调用`{SAVE_PLAN_TOOL_NAME}`工具来保存规划。
系统会自动为第一个处于PENDING状态的todo分配一个executor，executor执行该todo，然后将结果反馈给你。你需要根据反馈更新todo的状态、调整规划。
以上过程重复进行，直到你认为任务已经完成，然后你将所有todo标记为完成状态，并做一个总结。

## 不必过度拘泥于执行细节
你主要专注于把控（调整）规划的整体方向和进度，确保规划完整、可行，能够解决分配给团队的任务。
在设计/调整规划时，不必过度纠结于todo的执行细节，而是应该把这些细节交给executor去解决。因为你没有与任务环境交互的能力，只能空想，而executor可以边思考边尝试，在这方面要比你效率更高。

# executor

## executor的能力
executor可以执行的操作包括：
- 向其他executor, verifier提问
- 运行命令
- 与shell终端交互
- 读写文件
- 使用自身的视觉能力，读取、理解、处理部分多媒体文件（图片、pdf）
- 网页搜索和抓取
- 调用LLM api（支持多模态）进行批处理。llm调用成本较高，因此最好仅对抽样出的（或者可疑的）数据/结果进行分析）

## executor的记忆相互隔离
executor并不知道其他executor做了什么(除非其主动向其他executor提问)，所以你给executor的todo中最好包含执行该todo所依赖的信息。
如果缺乏相关信息，executor会自行调用工具去重新探索或者询问其它executor，但这样会导致效率降低，任务耗时增加，有任务超时的风险。

## executor可以向你提建议
executor对他自己的能力更为了解，同时他们工作在一线，获取到的信息更加具体。因此executor可能根据他的发现，向你建议更优的方案。
而你(planner)相比executor有更全局的任务视角优势，因此executor的建议也不一定正确，你需要综合任务历史和现状做出最优决策。

# 规划

## todo的粒度
- todo不宜过于简单，每个todo预期的executor执行步数应该至少为3,否则会增加与executor的协调成本，降低效率；
- todo也不宜过于复杂，否则executor可能因为过于复杂或者困难而无法在规定的步数内完成。

## 任务描述中可能会缺乏环境信息
在缺乏环境信息时，制定的规划可能会严重偏离任务原本的方向/目的，导致任务耗时增加，甚至任务失败。
因此当任务没有明确描述环境信息时（如当前目录下的文件、操作系统版本等），你需要先生成一个探索环境的规划，然后根据探索得到的信息，再生成解决任务的规划。

## 遵循最佳实践
收到任务后，要先分析任务所属的类别/领域，然后根据该领域的最佳实践来进行规划。
如：
任务：从xxx中恢复出yyy数据。
分析：该任务属于`数据恢复`领域，而数据恢复的最佳实践/准则是：在进行任何恢复操作之前，停止任何写入、立即进行备份。遵循该最佳实践可以防止操作意外导致数据损毁/丢失。
因此规划中的第1个todo必须是`数据备份`，其后才是`探索环境`等todo。

## 每项TODO完成后，要将其状态改为COMPLETED，而不是直接从plan中删除
由于每项todo是由不同的executor执行，删除todo会导致executor无法看到完整的规划，对其负责的todo感到困惑， 从而可能重复执行被，或者不遵守你的指令，自由发挥。
一旦出现这种局面，极易导致任务失败。
因此你在调整规划时，切记**保持规划的完整性，如果某项TODO是当前求解规划所必须的，且已经完成了，不要删除它**。
但如果你将规划推倒重来，或者大幅度调整了规划方向，导致某些TODO在新的规划中不再适用，那么你仍然可以自由地删除这些TODO。

{AGENT_FOLDER_INSTRUCTION}

{AGENT_COMMON_INSTRUCTION}
"""

logger = getLogger(__name__)


class TurnResult(BaseModel):
    ok: bool = True
    msg: str = ""

    @classmethod
    def success(cls, msg="") -> "TurnResult":
        return cls(ok=True, msg=msg)

    @classmethod
    def fail(cls, msg) -> "TurnResult":
        return cls(ok=False, msg=msg)


class PromptCacheInfo(BaseModel):
    index: int
    timestamp: float

    @classmethod
    def create(cls, index: int) -> "PromptCacheInfo":
        return cls(index=index, timestamp=time.perf_counter())

    def is_expired(self, duration: float) -> bool:
        now = time.perf_counter()
        return now - self.timestamp >= duration


class Planner(LLMAgent):
    def __init__(self):
        super().__init__()
        self._tool_registry = ToolRegistry(
            [Tool.from_callable(self._save_plan_tool, SAVE_PLAN_TOOL_NAME)]
        )
        self._tools = [
            tool_def.to_openai_format()
            for tool_def in self._tool_registry.get_tools_definition()
        ]

        self._messages: list[ModelInputMessage] = []
        self._prompt_cache_info: PromptCacheInfo | None = None
        self._task_history: TaskHistory = TaskHistory(team_member_results=[])
        self._failed_plans = list[Plan]()
        self._plan: Plan | None = None
        self._task_execution_summary = ""

        self._executor_id_counter = 0

    def next_executor_name(self) -> str:
        executor_name = get_executor_name(str(self._executor_id_counter))
        self._executor_id_counter += 1
        return executor_name

    async def _save_plan_tool(
        self,
        task_finished: Annotated[
            bool, Field(description="任务完成标记。在无需进一步规划时设置为为true。")
        ],
        todos: list[Todo],
    ) -> str:
        """必须提供完整的规划，完整是指没有缺少必要的步骤(todo)"""
        plan = Plan(task_finished=task_finished, todos=todos)
        self._plan = plan

        if not self._plan.task_finished:
            return "success."

        pending_todo_count = sum(
            todo.status == TodoStatus.PENDING for todo in self._plan.todos
        )
        if pending_todo_count != 0:
            raise RuntimeError(
                "错误：状态不一致，task_finished为true时，规划中不能包含处于PENDING状态的todo。"
            )
        return "success. " + "你已经将任务标记为完成，现在请生成任务总结。"

    @message_handler()
    async def start_task(self, task: Task, ctx: MessageContext) -> None:
        self._task = task
        planner_instruction = PLANNER_INSTRUCTION
        if task.env_notes:
            planner_instruction += task.env_notes
        prompt = "\n".join(
            [
                self._task.to_prompt(),
                self._task_config.to_prompt(),
            ]
        )
        self._messages = [
            {"role": "system", "content": planner_instruction},
            {"role": "user", "content": prompt},
        ]
        await self._publish_model_input(self._messages)
        await self._do_turn(ctx)

    @message_handler()
    async def receive_executor_result(
        self, executor_result: ExecutorResult, ctx: MessageContext
    ) -> None:
        self._task_history.team_member_results.append(executor_result)

        assert self._plan is not None
        finished_todo = self._plan.todos[executor_result.todo_index]
        content = (
            f"executor(名字：{executor_result.member_id.name})执行期间花费了{executor_result.step_count}步，调用了{executor_result.tool_call_count}次工具。\n"
            f"其执行了当前plan中的第{executor_result.todo_index + 1}项todo："
            f"<todo>{finished_todo.desc}。<todo>\n"
            f"其执行的最终状态是：{executor_result.end_reason.str_for_model()}\n"
        )
        if executor_result.extra_finished_todo_indexes:
            indexes_for_model = [
                index + 1 for index in executor_result.extra_finished_todo_indexes
            ]
            content += f"其还认为额外完成了第{indexes_for_model}项todo。\n"
        if executor_result.end_reason == ExecutionEndReason.STEPS_EXHAUSTED:
            content += (
                "注意：其因为超过了最大步数被中止了执行，被要求就取得的进展进行总结。\n"
            )
        content += f"executor的汇报（反馈）如下：<report>{executor_result.report_content}</report>\n"
        message: ModelInputMessage = {
            "role": "user",
            "content": content,
        }
        self._messages.append(message)
        await self._do_turn(ctx)

    async def _do_turn(self, ctx: MessageContext) -> None:
        LIMIT = 6
        for i in range(0, LIMIT):
            turn_result = await self._do_turn_inner(ctx)
            if turn_result.ok:
                return
            else:
                await self._debug_log(
                    "replanning {}/{}, reason: {}".format(i + 1, LIMIT, turn_result.msg)
                )
        await self._publish_output_event(
            TaskFailResult(content="planner stuck in do_turn loop")
        )

    async def _do_turn_inner(self, ctx: MessageContext) -> TurnResult:
        messages = copy(self._messages)
        if self._need_explicit_cache():
            last_message = copy(messages[-1])
            messages[-1] = last_message
            if isinstance(last_message, dict):
                ttl = None
                if self._plan:
                    team_member_result = self._task_history.team_member_results[-1]
                    if (
                        isinstance(team_member_result, ExecutorResult)
                        and team_member_result.end_reason
                        == ExecutionEndReason.COMPLETED
                    ):
                        pending_todo_count = sum(
                            todo.status == TodoStatus.PENDING
                            for todo in self._plan.todos
                        )
                        if len(self._plan.todos) > 1 and pending_todo_count == 1:
                            ttl = "5m"
                if ttl:
                    cache_control = {
                        "type": "ephemeral",
                        "ttl": ttl,
                    }
                    last_message["cache_control"] = cache_control  # pyright: ignore[reportGeneralTypeIssues]
                    self._prompt_cache_info = PromptCacheInfo.create(len(messages) - 1)

        choice = await self._acompletion(
            messages=messages,
            tools=self._tools,
        )
        tool_calls = choice.message.tool_calls

        if not tool_calls:
            self._messages.append(
                {
                    "role": "user",
                    "content": f"Ignore this message; it is only intended to prompt you to call the {SAVE_PLAN_TOOL_NAME} tool.",
                }
            )
            return TurnResult.fail("plan not generated")

        if len(tool_calls) > 1:
            return TurnResult.fail("multiple plan generated")

        tool_call = tool_calls[0]
        if tool_call.function.name is None:
            return TurnResult.fail("function name is empty")

        tool_exec_result = await self._tool_registry.call_tool(
            tool_call.function.name,
            tool_call.function.arguments,
        )
        self._messages.append(choice.message)
        self._messages.append(
            ChatCompletionToolMessageParam(
                content=convert_content_to_openai_format(tool_exec_result.content),
                tool_call_id=tool_call.id,
                role="tool",
            )
        )
        if not tool_exec_result.is_ok():
            return TurnResult.fail("save_plan tool call failed")

        if not self._plan or len(self._plan.todos) == 0:
            await self._publish_output_event(
                DebugTaskLog(content="planner: empty plan")
            )
            self._messages.append(
                {
                    "role": "user",
                    "content": "规划不能为空",
                }
            )
            return TurnResult.fail("plan is empty")

        for i, todo in enumerate(self._plan.todos):
            if todo.status == TodoStatus.PENDING:
                from judy_core.agents.executor import Executor

                await ctx.send(
                    Subtask(
                        task=self._task,
                        task_history=self._task_history.model_copy(deep=True),
                        plan=self._plan,
                        todo_index=i,
                    ),
                    ActorId(Executor.actor_kind, self.next_executor_name()),
                )
                return TurnResult.success("handle over to executor")
        if not self._plan.task_finished:
            return TurnResult.fail(
                "all todo of plan are done, but task_finished flag is not set"
            )

        choice = await self._gen_task_summary()
        assert choice.message.content is not None
        self._task_execution_summary = choice.message.content

        if self._task_config.disable_verifier:
            await self._publish_output_event(
                TaskSuccessResult(content=self._task_execution_summary)
            )
        else:
            await ctx.send(
                VerificationRequest(
                    task=self._task, task_execution_summary=self._task_execution_summary
                ),
                ActorId(Verifier.actor_kind, VERIFIER_NAME),
            )
        return TurnResult.success("task done")

    @message_handler()
    async def receive_verification_response(
        self, verification_result: VerificationResult, ctx: MessageContext
    ):
        self._task_history.team_member_results.append(verification_result)

        if verification_result.status == VerificationResultStatus.FAILED:
            # TODO 要求重新规划
            self._messages.append(
                {
                    "role": "user",
                    "content": f"任务验收没有通过，验收报告: <verify_result status={verification_result.status}>{verification_result.report_content}</verify_result>。",
                }
            )
            await self._do_turn(ctx)
        else:
            self._messages.append(
                {
                    "role": "user",
                    "content": f"任务验收通过，验收报告: <verify_result>{verification_result.report_content}</verify_result>。",
                }
            )
            await self._publish_output_event(
                TaskSuccessResult(content=self._task_execution_summary)
            )
        return

    async def _gen_task_summary(self) -> Choices:
        cache_control_injection_points = None
        if self._prompt_cache_info and not self._prompt_cache_info.is_expired(5 * 60):
            cache_control_injection_points = [
                {
                    "location": "message",
                    "index": self._prompt_cache_info.index,
                }
            ]
        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)
        return choice
