from itertools import chain

from judy_core.agents.prompts import AGENT_FOLDER_INSTRUCTION
from judy_core.agents.qa_agent import QAAgent
from judy_core.agents.react_agent import ReactAgent
from judy_core.events import CreateAgentFolder
from judy_core.messages import (
    Task,
    TeamMemberId,
    VerificationRequest,
    VerificationResult,
    VerificationResultStatus,
)
from judy_core.utils.tool import TOOL_RESULT_SUCCESS_CONTENT, Tool, ToolRegistry
from tongcell.core import ActorId, MessageContext, message_handler

FINISH_VERIFICATION_TOOL_NAME = "finish_verification"
VERIFIER_MAX_STEPS = 50

VERIFIER_STEPS_EXHAUSTED_PROMPT = """\
验证期间执行的步数已超限，中止验证。
"""

VERIFIER_FINISH_INSTRUCTION = f"""\
如果你认为已经可以出具验证结果和报告，你需要调用{FINISH_VERIFICATION_TOOL_NAME}工具来结束验证；否则你应该调用其它工具继续进行验证。\
"""

VERIFIER_REPORT_INSTRUCTION = """\
你需要就目前的发现和取得的进展进行汇报，汇报中可以包含你的建议。该汇报会被提交给planner，以便其及时调整规划。**你现在已经不能调用任何工具了，直接将汇报作为你的回复内容即可**。
"""

# TODO 添加构建测试用例的手段，如 黑盒测试 白盒测试
VERIFIER_MAIN_INSTRUCTION = f"""\
你(verifier)隶属于一个任务小组，该小组由一个planner、一个verifier和多个executor组成。
planner负责制定规划，executor负责执行规划中的各个todo，当planner认为任务已经完成后，会交给你(verifier)来验证/测试。
如果验证通过则任务完成，否则planner会根据你的反馈，调整规划，并分配executor来修正错误，修正之后会再次交给你来验证。

# 你的工作
你负责对根据任务要求，对交付的内容进行验证，并出具验证报告。
注意：
**禁止对交付的结果（代码或者数据文件）做任何修改。**
**不要尝试自行解决或者改进验证中发现的问题**。
你所发现的问题由planner和executor负责修复，他们有任务的完整历史信息和相关工具，在解决任务方面更有经验，更高效。


{AGENT_FOLDER_INSTRUCTION}
**作为verifier，交付目录对于你是只读的，禁止创建或者修改交付目录中的任何文件。你直接或者间接创建的文件（如编写的测试脚本、编译的产物等)，必须位于你的个人目录。**

# 测试case（验证case）
测试case来自于任务要求，因此在开始验证前，必须先分析任务中所提出的任务要求。注意：
- 当任务包含口语化表达或者非正式用语时，你必须敏锐地捕捉到其言外之意和隐性的要求。
- 当任务要求与环境冲突时，以任务要求为准。
详细内容见**任务要求**小结。
当所交付的解决方案不满足任务要求时，须判断为验收不通过。

## corner case
除非任务明确要求，你不必对非法输入等corner case进行验证，可以假定输入参数都是合法的参数。

### 测试脚本
如果任务提供了测试脚本，你必须运行该脚本来进行验证。
测试脚本可能需要特殊的运行方式才能得到正确的结果（比如用python解释器运行pytest脚本会导致进程正常退出，但实际上测试用例并没有被执行）。

### 测试数据集
有些问题可能是公开或者经典问题，这类问题通常会有公开的数据集。
对于此类问题，除了任务提供的数据集，你还可以通过下载额外的数据集来确保交付方案的覆盖率。

### 自定义验证方式
如果任务没有提供验证脚本，或者验证脚本的没有完全覆盖任务要求。那么你需要自行设计合理的验证方式。

### 交付结果的可重复性/幂等性
如果任务显式或者隐式要求（根据惯例或最佳实践）交付结果应该是可重复的/幂等的，应对交付结果的可重复性/幂等性进行验证。

# 你的两个阶段
## 验证阶段
在这个阶段，你需要根据任务要求，调用工具对交付的结果进行全面的验证。
{VERIFIER_FINISH_INSTRUCTION}

## 汇报阶段
在这个阶段，你需要就验证结果进行汇报，汇报内容会被提交给planner。

### 验证报告
当验证未通过时，为了帮助planner和executor定位问题和改进，你需要在验证报告中提供失败的case的复现方式（如验证脚本/数据的路径和运行方式）。

## 补充说明
一个任务可能会经历多次验证-修正循环，直至验证通过为止。因此你会在两个阶段间不断切换。

"""


# 你可以在调用{FINISH_VERIFICATION_TOOL_NAME}工具的同时，在tool call之外的回复内容中，同时给出你的汇报内容。
# 这样两个阶段就合并在了一起，可以减少一次交互延迟，提升效率。

CHECK_ITEM_ALL_TASK_REQUIREMENTS_AND_TEST_CASES_VERIFIED = """\
从任务描述中仔细分析出所有的：
- 需求
- 测试用例（应覆盖任务描述中提及的各类参数取值）
必须确保验证过程完整覆盖了以上内容。
"""

DEFAULT_VERIFIER_MANDATORY_CHECKLIST = {
    "all-task-requirements-and-test-cases-verified": CHECK_ITEM_ALL_TASK_REQUIREMENTS_AND_TEST_CASES_VERIFIED,
}

INVALID_CHECKLIST_MESSAGE_TEMPLATE = """\
tool call failed. 
when verification passed, following check items must be present in checklist.\
(Note that the check item names are connected with hyphens `-`, not underscores `_`.): 
{}.
If you believe these check items have already been verified, please pass them in. \
If you forgot to verify them, you may continue with your verification phase \
and call this tool after the verification has been completed successfully.
"""


class Verifier(QAAgent, ReactAgent):
    def __init__(self):
        super().__init__()
        self._mandatory_checklist: dict[str, str] = DEFAULT_VERIFIER_MANDATORY_CHECKLIST
        self._verification_result_status: VerificationResultStatus | None = None
        self._prev_checklist_arg_invalid_count = 0

    def _react_exit_prompt(self) -> str:
        return VERIFIER_FINISH_INSTRUCTION

    async def _on_turn_done(self, ctx: MessageContext) -> None:
        if self._verification_result_status is None:
            self._step_count += 1
            if self._step_count < VERIFIER_MAX_STEPS:
                await self._do_turn(ctx)
            else:
                await self._debug_log(f"exceeded max steps: {VERIFIER_MAX_STEPS}")
                self._verification_result_status = (
                    VerificationResultStatus.STEPS_EXHAUSTED
                )
                self._messages.append(
                    {
                        "role": "user",
                        "content": VERIFIER_STEPS_EXHAUSTED_PROMPT,
                    }
                )
                report_content = await self._gen_report(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(ctx)

            await self._send_report(report_content, ctx)

    async def finish_verification(
        self,
        verification_result_status: VerificationResultStatus,
        checklist: list[str] | None = None,
    ) -> str:
        """结束验证并给出验证结果。"""
        if verification_result_status == VerificationResultStatus.PASSED:
            if checklist is None:
                checklist = []
            missing_checklist_items = []
            for check_item in self._mandatory_checklist.keys():
                if check_item not in checklist:
                    missing_checklist_items.append(check_item)
            if len(missing_checklist_items) > 0:
                self._prev_checklist_arg_invalid_count += 1
                if self._prev_checklist_arg_invalid_count <= 3:
                    raise RuntimeError(
                        INVALID_CHECKLIST_MESSAGE_TEMPLATE.format(
                            missing_checklist_items
                        )
                    )
                else:
                    await self._debug_log(
                        f"The LLM failed to correct the checklist parameter in the finish_verification call. To prevent an infinite loop, the verification step was allowed to proceed. {missing_checklist_items}"
                    )
        self._verification_result_status = verification_result_status
        self._prev_checklist_arg_invalid_count = 0
        return TOOL_RESULT_SUCCESS_CONTENT + ". " + VERIFIER_REPORT_INSTRUCTION

    def _get_finish_verification_tool_description(self) -> str:
        description_segs = ["结束验证并给出验证结果。"]
        if self._mandatory_checklist:
            description_segs.append(
                "在给出验证通过的结论前，必须完成以下check items的检查和验证，如果验证没有问题，则将这些check items添加到checklist参数中。"
            )
            for (
                check_item,
                check_item_desc,
            ) in self._mandatory_checklist.items():
                description_segs.append(
                    f"## check_item: `{check_item}`\n{check_item_desc}"
                )
        description = "\n".join(description_segs)
        return description

    async def _init_verification(self, task: Task) -> None:
        self._verification_result_status = None

        self._internal_tool_registry = ToolRegistry(
            [
                Tool.from_callable(
                    self.finish_verification,
                    FINISH_VERIFICATION_TOOL_NAME,
                    self._get_finish_verification_tool_description(),
                ),
            ],
        )
        self._tools = [
            tool_def.to_openai_format()
            for tool_def in chain(
                self._model_config.tools_definition,
                self._internal_tool_registry.get_tools_definition(),
            )
        ]

        verifier_instruction = VERIFIER_MAIN_INSTRUCTION
        if task.env_notes:
            verifier_instruction += task.env_notes

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

        prompt = "\n".join(
            [
                task.to_prompt(),
                self._task_config.to_prompt(personal_dir),
            ]
        )

        self._messages = [
            {"role": "system", "content": verifier_instruction},
            {"role": "user", "content": prompt},
        ]

        await self._publish_model_input(self._messages)

    @message_handler()
    async def start_verification(
        self, verification_request: VerificationRequest, ctx: MessageContext
    ) -> None:
        if len(self._messages) == 0:
            if verification_request.task.verifier_mandatory_checklist:
                self._mandatory_checklist.update(
                    verification_request.task.verifier_mandatory_checklist
                )
            await self._init_verification(verification_request.task)
        self._step_count = 0
        self._verification_result_status = None

        self._messages.append(
            {
                "role": "user",
                "content": f"planner认为当前任务已经完成，现在你需要进行验证。planner给出的task_summary如下: <task_summary>{verification_request.task_execution_summary}</task_summary>",
            }
        )

        await self._do_turn(ctx)

    async def _gen_report(self, ctx: MessageContext) -> str:
        choice = await self._acompletion_without_tool_calls(
            messages=self._messages,
            tools=self._tools,
            cache_control_injection_points=[
                {
                    "location": "message",
                    "index": -3,
                }
            ],
        )
        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._verification_result_status is not None

        await ctx.send(
            VerificationResult(
                member_id=TeamMemberId.from_actor_id(self._id),
                status=self._verification_result_status,
                report_content=report_content,
            ),
            ActorId.default(Planner.actor_kind),
        )
