"""子智能体专用终端，实现工具白名单与完成工具。"""

import json
from pathlib import Path
from typing import Dict, Optional, List

from core.web_terminal import WebTerminal
from config import PROMPTS_DIR

FORBIDDEN_SUB_AGENT_TOOLS = {"update_memory", "todo_create", "todo_update_task", "todo_finish", "todo_finish_confirm"}


class SubAgentTerminal(WebTerminal):
    """子智能体 Web 终端，限制工具并提供 finish_sub_agent。"""

    def __init__(
        self,
        *,
        workspace_dir: str,
        data_dir: str,
        metadata: Dict,
        message_callback=None,
    ):
        super().__init__(
            project_path=workspace_dir,
            thinking_mode=True,
            message_callback=message_callback,
            data_dir=data_dir,
        )
        self.sub_agent_meta = metadata
        self.finish_callback = None
        self._system_prompt_cache: Optional[str] = None

    def set_finish_callback(self, callback):
        self.finish_callback = callback

    def load_prompt(self, name: str) -> str:
        if name != "main_system":
            return super().load_prompt(name)

        if self._system_prompt_cache:
            return self._system_prompt_cache

        template_path = Path(PROMPTS_DIR) / "sub_agent_system.txt"
        if not template_path.exists():
            return super().load_prompt(name)

        template = template_path.read_text(encoding="utf-8")
        data = {
            "summary": self.sub_agent_meta.get("summary", ""),
            "task": self.sub_agent_meta.get("task", ""),
            "workspace": self.sub_agent_meta.get("workspace_dir", ""),
            "references": self.sub_agent_meta.get("references_dir", ""),
            "deliverables": self.sub_agent_meta.get("deliverables_dir", ""),
            "target_project_dir": self.sub_agent_meta.get("target_project_dir", ""),
            "agent_id": self.sub_agent_meta.get("agent_id", ""),
            "task_id": self.sub_agent_meta.get("task_id", ""),
        }
        self._system_prompt_cache = template.format(**data)
        return self._system_prompt_cache

    def define_tools(self) -> List[Dict]:
        tools = super().define_tools()
        filtered: List[Dict] = []
        for tool in tools:
            name = tool.get("function", {}).get("name")
            if name in FORBIDDEN_SUB_AGENT_TOOLS:
                continue
            filtered.append(tool)

        filtered.append({
            "type": "function",
            "function": {
                "name": "finish_sub_agent",
                "description": (
                    "当你确定交付成果已准备完毕时调用此工具。调用前请确认 deliverables 文件夹存在 result.md，"
                    "其中包含交付说明。参数 reason 用于向主智能体总结本轮完成情况。"
                ),
                "parameters": {
                    "type": "object",
                    "properties": {
                        "reason": {
                            "type": "string",
                            "description": "向主智能体说明任务完成情况、交付内容、下一步建议。"
                        }
                    },
                    "required": ["reason"]
                }
            }
        })
        return filtered

    async def handle_tool_call(self, tool_name: str, arguments: Dict) -> str:
        if tool_name == "finish_sub_agent":
            result = self._finalize_sub_agent(arguments or {})
            return json.dumps(result, ensure_ascii=False)
        return await super().handle_tool_call(tool_name, arguments)

    def _finalize_sub_agent(self, arguments: Dict) -> Dict:
        deliverables_dir = Path(self.sub_agent_meta.get("deliverables_dir", self.project_path))
        result_md = deliverables_dir / "result.md"
        if not result_md.exists():
            return {
                "success": False,
                "error": "deliverables 目录缺少 result.md，无法结束任务。"
            }
        content = result_md.read_text(encoding="utf-8").strip()
        if not content:
            return {
                "success": False,
                "error": "result.md 为空，请写入任务总结与交付说明后再结束任务。"
            }
        reason = (arguments.get("reason") or "").strip()
        if not reason:
            return {
                "success": False,
                "error": "缺少 reason 字段，请说明完成情况。"
            }
        result = {
            "success": True,
            "message": "子智能体任务已标记为完成。",
            "reason": reason,
        }
        if self.finish_callback:
            try:
                self.finish_callback(result)
            except Exception:
                pass
        return result