ForcePilot/backend/package/yuxi/channel/extensions/zalouser/tools.py

from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)

AGENT_TOOL_NAME = "zalouser"

AGENT_TOOL_SCHEMA = {
    "name": AGENT_TOOL_NAME,
    "description": (
        "Zalo Personal actions: send messages/images/links/stickers, list friends/groups, "
        "manage groups (create/disband/add/remove/rename/transfer/deputies), "
        "polls, find users, manage friends (block/unblock/alias/requests), "
        "notes, cards, pin conversations, export cookie, get profile."
    ),
    "parameters": {
        "type": "object",
        "properties": {
            "action": {
                "type": "string",
                "enum": [
                    "send", "image", "link", "delete",
                    "poll_create", "poll_lock",
                    "friends", "groups", "group_members",
                    "me", "status",
                    "find_user", "user_info", "group_info",
                    "create_group", "disband_group",
                    "add_member", "remove_member",
                    "rename_group", "change_avatar", "transfer_owner",
                    "add_deputy", "remove_deputy",
                    "block", "unblock", "change_alias",
                    "send_friend_request", "accept_friend_request",
                    "send_sticker", "send_card",
                    "create_note", "edit_note",
                    "pin", "cookie",
                ],
                "description": "Action to perform",
            },
            "thread_id": {
                "type": "string",
                "description": "Target thread/chat ID",
            },
            "message": {
                "type": "string",
                "description": "Message text or caption to send",
            },
            "is_group": {
                "type": "boolean",
                "description": "Whether the target is a group chat",
            },
            "profile": {
                "type": "string",
                "description": "Profile name for multi-account setups",
            },
            "query": {
                "type": "string",
                "description": "Search query for friends/groups listing",
            },
            "url": {
                "type": "string",
                "description": "Media or link URL",
            },
            "message_id": {
                "type": "string",
                "description": "ID of the message to delete (required for delete action)",
            },
            "question": {
                "type": "string",
                "description": "Poll question",
            },
            "options": {
                "type": "array",
                "items": {"type": "string"},
                "description": "Poll options (list of strings)",
            },
            "poll_id": {
                "type": "string",
                "description": "Poll ID to lock",
            },
            "phone_number": {
                "type": "string",
                "description": "Phone number to find user",
            },
            "user_id": {
                "type": "string",
                "description": "Zalo user ID",
            },
            "group_id": {
                "type": "string",
                "description": "Zalo group ID",
            },
            "name": {
                "type": "string",
                "description": "Name (group name, contact name, etc.)",
            },
            "member_ids": {
                "type": "array",
                "items": {"type": "string"},
                "description": "List of member/thread IDs",
            },
            "alias": {
                "type": "string",
                "description": "Friend alias/remark name",
            },
            "avatar_url": {
                "type": "string",
                "description": "URL of the group avatar image",
            },
            "member_id": {
                "type": "string",
                "description": "Member ID to transfer ownership to",
            },
            "msg": {
                "type": "string",
                "description": "Message for friend request",
            },
            "sticker_obj": {
                "type": "object",
                "description": "Sticker object metadata",
            },
            "phone": {
                "type": "string",
                "description": "Phone number for contact card",
            },
            "title": {
                "type": "string",
                "description": "Note title",
            },
            "content": {
                "type": "string",
                "description": "Note content",
            },
            "color": {
                "type": "integer",
                "description": "Note color (0=default)",
            },
            "note_id": {
                "type": "string",
                "description": "Note ID to edit",
            },
            "thread_ids": {
                "type": "array",
                "items": {"type": "string"},
                "description": "List of thread IDs to pin",
            },
        },
        "required": ["action"],
    },
}


class ZaloUserAgentTools:
    def __init__(self):
        self._outbound = None
        self._directory = None

    def set_auth(self, auth):
        self._auth = auth

    def set_outbound(self, outbound):
        self._outbound = outbound

    def set_directory(self, directory):
        self._directory = directory

    def get_tool_def(self) -> dict:
        return AGENT_TOOL_SCHEMA

    async def execute(self, params: dict, context: dict | None = None) -> dict[str, Any]:
        action = params.get("action", "")
        if not action:
            return {"success": False, "error": "action is required"}

        thread_id = params.get("thread_id", "")
        message = params.get("message", "")
        url = params.get("url", "")
        is_group = params.get("is_group", False)

        if action == "send":
            if not thread_id or not message:
                return {"success": False, "error": "thread_id and message required"}
            await self._send_text(thread_id, message, is_group)
            return {"success": True, "action": "send", "thread_id": thread_id}

        if action == "image":
            if not thread_id or not url:
                return {"success": False, "error": "thread_id and url required"}
            await self._send_image(thread_id, url, message, is_group)
            return {"success": True, "action": "image", "thread_id": thread_id}

        if action == "link":
            if not thread_id or not url:
                return {"success": False, "error": "thread_id and url required"}
            await self._send_link(thread_id, url, message, is_group)
            return {"success": True, "action": "link", "thread_id": thread_id}

        if action == "delete":
            message_id = params.get("message_id", "")
            if not thread_id or not message_id:
                return {"success": False, "error": "thread_id and message_id required"}
            await self._delete_message(thread_id, message_id)
            return {"success": True, "action": "delete", "thread_id": thread_id}

        if action == "poll_create":
            question = params.get("question", "")
            options = params.get("options", [])
            if not thread_id or not question or not options:
                return {"success": False, "error": "thread_id, question, and options required"}
            result = await self._create_poll(thread_id, question, options)
            return {"success": True, "action": "poll_create", "result": result}

        if action == "poll_lock":
            poll_id = params.get("poll_id", "")
            if not thread_id or not poll_id:
                return {"success": False, "error": "thread_id and poll_id required"}
            result = await self._lock_poll(thread_id, poll_id)
            return {"success": True, "action": "poll_lock", "result": result}

        if action in ("friends", "groups"):
            query = params.get("query", "")
            if action == "friends":
                result = await self._list_friends(query)
            else:
                result = await self._list_groups(query)
            return {"success": True, "action": action, "result": result}

        if action == "group_members":
            group_id = params.get("group_id", "")
            if not group_id:
                return {"success": False, "error": "group_id required"}
            result = await self._list_group_members(group_id)
            return {"success": True, "action": "group_members", "result": result}

        if action == "me":
            result = await self._get_me()
            return {"success": True, "action": "me", "result": result}

        if action == "status":
            result = await self._check_status()
            return {"success": True, "action": "status", "result": result}

        if action == "find_user":
            phone_number = params.get("phone_number", "")
            if not phone_number:
                return {"success": False, "error": "phone_number required"}
            result = await self._find_user(phone_number)
            return {"success": True, "action": "find_user", "result": result}

        if action == "user_info":
            user_id = params.get("user_id", "")
            if not user_id:
                return {"success": False, "error": "user_id required"}
            result = await self._get_user_info(user_id)
            return {"success": True, "action": "user_info", "result": result}

        if action == "group_info":
            group_id = params.get("group_id", "")
            if not group_id:
                return {"success": False, "error": "group_id required"}
            result = await self._get_group_info(group_id)
            return {"success": True, "action": "group_info", "result": result}

        group_id = params.get("group_id", "")
        name = params.get("name", "")

        if action == "create_group":
            member_ids = params.get("member_ids", [])
            if not name or not member_ids:
                return {"success": False, "error": "name and member_ids required"}
            result = await self._create_group(name, member_ids)
            return {"success": True, "action": "create_group", "result": result}

        if action == "disband_group":
            if not group_id:
                return {"success": False, "error": "group_id required"}
            result = await self._disband_group(group_id)
            return {"success": True, "action": "disband_group", "result": result}

        user_id = params.get("user_id", "")

        if action == "add_member":
            if not user_id or not group_id:
                return {"success": False, "error": "user_id and group_id required"}
            result = await self._add_member(user_id, group_id)
            return {"success": True, "action": "add_member", "result": result}

        if action == "remove_member":
            if not user_id or not group_id:
                return {"success": False, "error": "user_id and group_id required"}
            result = await self._remove_member(user_id, group_id)
            return {"success": True, "action": "remove_member", "result": result}

        if action == "rename_group":
            if not group_id or not name:
                return {"success": False, "error": "group_id and name required"}
            result = await self._rename_group(group_id, name)
            return {"success": True, "action": "rename_group", "result": result}

        avatar_url = params.get("avatar_url", "")
        if action == "change_avatar":
            if not group_id or not avatar_url:
                return {"success": False, "error": "group_id and avatar_url required"}
            result = await self._change_avatar(group_id, avatar_url)
            return {"success": True, "action": "change_avatar", "result": result}

        member_id = params.get("member_id", "")
        if action == "transfer_owner":
            if not group_id or not member_id:
                return {"success": False, "error": "group_id and member_id required"}
            result = await self._transfer_owner(group_id, member_id)
            return {"success": True, "action": "transfer_owner", "result": result}

        member_ids = params.get("member_ids", [])
        if action == "add_deputy":
            if not group_id or not member_ids:
                return {"success": False, "error": "group_id and member_ids required"}
            result = await self._add_deputy(group_id, member_ids)
            return {"success": True, "action": "add_deputy", "result": result}

        if action == "remove_deputy":
            if not group_id or not member_ids:
                return {"success": False, "error": "group_id and member_ids required"}
            result = await self._remove_deputy(group_id, member_ids)
            return {"success": True, "action": "remove_deputy", "result": result}

        if action == "block":
            if not user_id:
                return {"success": False, "error": "user_id required"}
            result = await self._block_user(user_id)
            return {"success": True, "action": "block", "result": result}

        if action == "unblock":
            if not user_id:
                return {"success": False, "error": "user_id required"}
            result = await self._unblock_user(user_id)
            return {"success": True, "action": "unblock", "result": result}

        alias = params.get("alias", "")
        if action == "change_alias":
            if not user_id or not alias:
                return {"success": False, "error": "user_id and alias required"}
            result = await self._change_alias(user_id, alias)
            return {"success": True, "action": "change_alias", "result": result}

        msg = params.get("msg", "")
        if action == "send_friend_request":
            if not user_id:
                return {"success": False, "error": "user_id required"}
            result = await self._send_friend_request(user_id, msg)
            return {"success": True, "action": "send_friend_request", "result": result}

        if action == "accept_friend_request":
            if not user_id:
                return {"success": False, "error": "user_id required"}
            result = await self._accept_friend_request(user_id)
            return {"success": True, "action": "accept_friend_request", "result": result}

        if action == "send_sticker":
            sticker_obj = params.get("sticker_obj", {})
            if not thread_id or not sticker_obj:
                return {"success": False, "error": "thread_id and sticker_obj required"}
            result = await self._send_sticker(thread_id, sticker_obj)
            return {"success": True, "action": "send_sticker", "result": result}

        phone = params.get("phone", "")
        contact_name = params.get("name", "")
        if action == "send_card":
            if not thread_id or not user_id:
                return {"success": False, "error": "thread_id and user_id required"}
            result = await self._send_card(thread_id, user_id, phone, contact_name)
            return {"success": True, "action": "send_card", "result": result}

        title = params.get("title", "")
        content = params.get("content", "")
        color = params.get("color", 0)
        if action == "create_note":
            if not thread_id or not title or not content:
                return {"success": False, "error": "thread_id, title, and content required"}
            result = await self._create_note(thread_id, title, content, color)
            return {"success": True, "action": "create_note", "result": result}

        note_id = params.get("note_id", "")
        if action == "edit_note":
            if not thread_id or not note_id or not title or not content:
                return {"success": False, "error": "thread_id, note_id, title, and content required"}
            result = await self._edit_note(thread_id, note_id, title, content, color)
            return {"success": True, "action": "edit_note", "result": result}

        thread_ids = params.get("thread_ids", [])
        if action == "pin":
            if not thread_ids:
                return {"success": False, "error": "thread_ids required"}
            result = await self._pin(thread_ids)
            return {"success": True, "action": "pin", "result": result}

        if action == "cookie":
            result = await self._get_cookie()
            return {"success": True, "action": "cookie", "result": result}

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

    def execute_sync(self, params: dict, context: dict | None = None) -> dict[str, Any]:
        action = params.get("action", "")
        profile = params.get("profile", "default")

        if not action:
            return {"success": False, "error": "action is required"}

        read_only_sync = {
            "status": {"authenticated": True, "profile": profile},
            "friends": [],
            "groups": [],
            "group_members": [],
            "me": None,
        }

        if action in read_only_sync:
            return {"success": True, "action": action, "result": read_only_sync[action]}

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

    async def _send_text(self, thread_id: str, message: str, is_group: bool = False) -> None:
        if self._outbound:
            await self._outbound.send_text(thread_id, message)

    async def _send_image(self, thread_id: str, url: str, caption: str = "", is_group: bool = False) -> None:
        if self._outbound:
            await self._outbound.send_media(thread_id, url, "image")

    async def _send_link(self, thread_id: str, url: str, caption: str = "", is_group: bool = False) -> None:
        if self._outbound:
            await self._outbound.send_link(thread_id, url, caption)

    async def _delete_message(self, thread_id: str, message_id: str) -> None:
        if self._outbound:
            await self._outbound.delete_message(thread_id, message_id)

    async def _list_friends(self, query: str = "") -> list[dict]:
        if self._directory:
            friends = await self._directory.list_friends(query)
            return [{"user_id": f.user_id, "display_name": f.display_name} for f in friends]
        return []

    async def _list_groups(self, query: str = "") -> list[dict]:
        if self._directory:
            groups = await self._directory.list_groups(query)
            return [{"group_id": g.group_id, "name": g.name, "member_count": g.member_count} for g in groups]
        return []

    async def _list_group_members(self, group_id: str) -> list[dict]:
        if self._directory:
            members = await self._directory.list_group_members(group_id)
            return [{"user_id": m.user_id, "display_name": m.display_name, "is_admin": m.is_admin} for m in members]
        return []

    async def _get_me(self) -> dict | None:
        if self._directory:
            info = await self._directory.get_self()
            if info:
                return {"user_id": info.user_id, "display_name": info.display_name}
        return None

    async def _check_status(self) -> dict:
        return {"authenticated": True}

    async def _create_poll(self, group_id: str, question: str, options: list[str]) -> dict:
        if self._outbound:
            return await self._outbound.send_poll(group_id, question, options)
        return {"success": False, "error": "outbound not set"}

    async def _lock_poll(self, group_id: str, poll_id: str) -> dict:
        if self._outbound:
            return await self._outbound.end_poll(group_id, poll_id)
        return {"success": False, "error": "outbound not set"}

    async def _find_user(self, phone_number: str) -> dict | None:
        if self._directory:
            return await self._directory.find_user(phone_number)
        return None

    async def _get_user_info(self, user_id: str) -> dict | None:
        if self._directory:
            info = await self._directory.get_user_info(user_id)
            if info:
                return {"user_id": info.user_id, "display_name": info.display_name}
        return None

    async def _get_group_info(self, group_id: str) -> dict | None:
        if self._directory:
            info = await self._directory.get_group_info(group_id)
            if info:
                return {"group_id": info.group_id, "name": info.name, "member_count": info.member_count}
        return None

    async def _create_group(self, name: str, member_ids: list[str]) -> dict:
        if self._directory:
            return await self._directory.create_group(name, member_ids)
        return {"success": False, "error": "directory not set"}

    async def _disband_group(self, group_id: str) -> dict:
        if self._directory:
            return await self._directory.disband_group(group_id)
        return {"success": False, "error": "directory not set"}

    async def _add_member(self, user_id: str, group_id: str) -> dict:
        if self._directory:
            return await self._directory.add_user_to_group(user_id, group_id)
        return {"success": False, "error": "directory not set"}

    async def _remove_member(self, user_id: str, group_id: str) -> dict:
        if self._directory:
            return await self._directory.remove_user_from_group(user_id, group_id)
        return {"success": False, "error": "directory not set"}

    async def _rename_group(self, group_id: str, name: str) -> dict:
        if self._directory:
            return await self._directory.change_group_name(group_id, name)
        return {"success": False, "error": "directory not set"}

    async def _change_avatar(self, group_id: str, avatar_url: str) -> dict:
        if self._directory:
            return await self._directory.change_group_avatar(group_id, avatar_url)
        return {"success": False, "error": "directory not set"}

    async def _transfer_owner(self, group_id: str, member_id: str) -> dict:
        if self._directory:
            return await self._directory.change_group_owner(group_id, member_id)
        return {"success": False, "error": "directory not set"}

    async def _add_deputy(self, group_id: str, member_ids: list[str]) -> dict:
        if self._directory:
            return await self._directory.add_group_deputies(group_id, member_ids)
        return {"success": False, "error": "directory not set"}

    async def _remove_deputy(self, group_id: str, member_ids: list[str]) -> dict:
        if self._directory:
            return await self._directory.remove_group_deputies(group_id, member_ids)
        return {"success": False, "error": "directory not set"}

    async def _block_user(self, user_id: str) -> dict:
        if self._directory:
            return await self._directory.block_user(user_id)
        return {"success": False, "error": "directory not set"}

    async def _unblock_user(self, user_id: str) -> dict:
        if self._directory:
            return await self._directory.unblock_user(user_id)
        return {"success": False, "error": "directory not set"}

    async def _change_alias(self, user_id: str, alias: str) -> dict:
        if self._directory:
            return await self._directory.change_friend_alias(user_id, alias)
        return {"success": False, "error": "directory not set"}

    async def _send_friend_request(self, user_id: str, msg: str = "") -> dict:
        if self._directory:
            return await self._directory.send_friend_request(user_id, msg)
        return {"success": False, "error": "directory not set"}

    async def _accept_friend_request(self, user_id: str) -> dict:
        if self._directory:
            return await self._directory.accept_friend_request(user_id)
        return {"success": False, "error": "directory not set"}

    async def _send_sticker(self, thread_id: str, sticker_obj: dict) -> dict:
        if self._outbound:
            return await self._outbound.send_sticker(thread_id, sticker_obj)
        return {"success": False, "error": "outbound not set"}

    async def _send_card(self, thread_id: str, user_id: str, phone: str, name: str) -> dict:
        if self._outbound:
            return await self._outbound.send_card(thread_id, user_id, phone, name)
        return {"success": False, "error": "outbound not set"}

    async def _create_note(self, group_id: str, title: str, content: str, color: int = 0) -> dict:
        if self._outbound:
            return await self._outbound.create_note(group_id, title, content, color)
        return {"success": False, "error": "outbound not set"}

    async def _edit_note(self, group_id: str, note_id: str, title: str, content: str, color: int = 0) -> dict:
        if self._outbound:
            return await self._outbound.edit_note(group_id, note_id, title, content, color)
        return {"success": False, "error": "outbound not set"}

    async def _pin(self, thread_ids: list[str]) -> dict:
        if self._directory:
            return await self._directory.pin_conversations(thread_ids)
        return {"success": False, "error": "directory not set"}

    async def _get_cookie(self) -> dict:
        if self._auth:
            return await self._auth.get_cookie()
        return {"success": False, "error": "auth adapter not available"}