ForcePilot/backend/package/yuxi/channels/adapters/imessage/agent_tools.py

from __future__ import annotations

from typing import Any


def describe_imessage_pairing_tool() -> dict:
    return {
        "type": "function",
        "function": {
            "name": "imessage_manage_pairing",
            "description": "管理 iMessage 配对请求：列出待审批的用户、批准或拒绝配对",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["list", "approve", "reject"],
                        "description": "操作类型：list=列出待审批, approve=批准配对, reject=拒绝配对",
                    },
                    "code": {
                        "type": "string",
                        "description": "配对码（approve/reject 操作需要）",
                    },
                },
                "required": ["action"],
            },
        },
    }


def describe_imessage_allowlist_tool() -> dict:
    return {
        "type": "function",
        "function": {
            "name": "imessage_manage_allowlist",
            "description": "管理 iMessage 白名单：添加或移除联系人/群组白名单",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["add_dm", "remove_dm", "add_group", "remove_group", "list"],
                        "description": "操作类型",
                    },
                    "handle": {
                        "type": "string",
                        "description": "电话号码或 chat_guid",
                    },
                },
                "required": ["action"],
            },
        },
    }


def describe_imessage_security_status_tool() -> dict:
    return {
        "type": "function",
        "function": {
            "name": "imessage_security_status",
            "description": "查看 iMessage 渠道的安全状态和警告信息",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": [],
            },
        },
    }


class IMessageAgentToolFactory:
    def __init__(self, adapter: Any):
        self._adapter = adapter

    def list_tools(self) -> list[dict]:
        return [
            describe_imessage_pairing_tool(),
            describe_imessage_allowlist_tool(),
            describe_imessage_security_status_tool(),
        ]

    async def execute_tool(self, tool_name: str, params: dict) -> Any:
        if tool_name == "imessage_manage_pairing":
            return await self._manage_pairing(params)
        if tool_name == "imessage_manage_allowlist":
            return await self._manage_allowlist(params)
        if tool_name == "imessage_security_status":
            return self._security_status()
        raise ValueError(f"Unknown tool: {tool_name}")

    async def _manage_pairing(self, params: dict) -> dict:
        action = params.get("action", "list")

        if action == "list":
            pending = self._adapter.list_pending_pairings()
            paired = self._adapter._pairing.get_paired_handles()
            return {"pending": pending, "paired": paired}

        code = params.get("code", "")
        if action == "approve":
            ok = self._adapter.approve_pairing(code)
            if ok:
                await self._adapter._pairing.save_allowlist()
            return {"success": ok, "message": "Pairing approved" if ok else "Invalid or expired code"}

        if action == "reject":
            ok = self._adapter.reject_pairing(code)
            return {"success": ok, "message": "Pairing rejected" if ok else "Invalid or expired code"}

        return {"success": False, "error": f"Unknown action: {action}"}

    async def _manage_allowlist(self, params: dict) -> dict:
        action = params.get("action", "list")

        if action == "list":
            return {
                "dm_allowlist": self._adapter._security.allow_list,
                "group_allowlist": self._adapter._security.group_allow_list,
            }

        handle = params.get("handle", "")
        if not handle:
            return {"success": False, "error": "handle is required"}

        if action == "add_dm":
            self._adapter._security.add_to_allow_list(handle)
            return {"success": True, "message": f"Added {handle} to DM allowlist"}

        if action == "remove_dm":
            ok = self._adapter._security.remove_from_allow_list(handle)
            return {"success": ok, "message": f"{'Removed' if ok else 'Not found'}: {handle}"}

        if action == "add_group":
            self._adapter._security.add_to_group_allow_list(handle)
            return {"success": True, "message": f"Added {handle} to group allowlist"}

        if action == "remove_group":
            ok = self._adapter._security.remove_from_group_allow_list(handle)
            return {"success": ok, "message": f"{'Removed' if ok else 'Not found'}: {handle}"}

        return {"success": False, "error": f"Unknown action: {action}"}

    def _security_status(self) -> dict:
        warnings = self._adapter.get_security_warnings()
        return {
            "dm_policy": self._adapter._security.dm_policy.value,
            "group_policy": self._adapter._security.group_policy.value,
            "require_mention": self._adapter._security.require_mention,
            "warnings": warnings,
            "paired_count": len(self._adapter._pairing.get_paired_handles()),
        }