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": [],
            },
        },
    }


def describe_imessage_poll_tool() -> dict:
    return {
        "type": "function",
        "function": {
            "name": "imessage_manage_polls",
            "description": "管理 iMessage 投票：创建、关闭、查看投票",
            "parameters": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": ["create", "close", "list", "get"],
                        "description": "操作类型",
                    },
                    "chat_guid": {
                        "type": "string",
                        "description": "聊天 ID（创建/列出时需要）",
                    },
                    "question": {
                        "type": "string",
                        "description": "投票问题（创建时需要）",
                    },
                    "options": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "选项列表（创建时需要）",
                    },
                    "poll_id": {
                        "type": "string",
                        "description": "投票 ID（关闭/查看时需要）",
                    },
                    "duration_minutes": {
                        "type": "integer",
                        "description": "投票持续时间（分钟），默认 5 分钟",
                    },
                },
                "required": ["action"],
            },
        },
    }


def describe_imessage_audit_tool() -> dict:
    return {
        "type": "function",
        "function": {
            "name": "imessage_audit_log",
            "description": "查询 iMessage 渠道审计日志",
            "parameters": {
                "type": "object",
                "properties": {
                    "limit": {
                        "type": "integer",
                        "description": "返回记录数上限，默认 20",
                    },
                },
                "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(),
            describe_imessage_poll_tool(),
            describe_imessage_audit_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()
        if tool_name == "imessage_manage_polls":
            return await self._manage_polls(params)
        if tool_name == "imessage_audit_log":
            return self._audit_log(params)
        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()),
        }

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

        if action == "list":
            active = [
                {
                    "poll_id": p.poll_id,
                    "chat_guid": p.chat_guid,
                    "question": p.question,
                    "closed": p.closed,
                }
                for p in self._adapter._poll_manager._polls.values()
                if not p.closed
            ]
            return {"active_polls": active}

        if action == "create":
            chat_guid = params.get("chat_guid", "")
            question = params.get("question", "")
            options = params.get("options", [])
            duration_minutes = params.get("duration_minutes", 5)
            if not chat_guid or not question or not options:
                return {"success": False, "error": "chat_guid, question, and options are required"}
            result = await self._adapter.create_poll(
                chat_id=chat_guid,
                question=question,
                options=options,
                duration_minutes=duration_minutes,
            )
            return {"success": True, "poll": result}

        if action == "close":
            poll_id = params.get("poll_id", "")
            if not poll_id:
                return {"success": False, "error": "poll_id is required"}
            result = await self._adapter.close_poll("", poll_id)
            return {"success": True, "result": result}

        if action == "get":
            poll_id = params.get("poll_id", "")
            if not poll_id:
                return {"success": False, "error": "poll_id is required"}
            result = self._adapter.get_poll(poll_id)
            if result is None:
                return {"success": False, "error": f"Poll not found: {poll_id}"}
            return {"success": True, "poll": result}

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

    def _audit_log(self, params: dict) -> dict:
        limit = params.get("limit", 20)
        records = self._adapter._audit_logger.get_recent(limit=limit)
        return {"records": records, "count": len(records)}