from typing import List, Dict

from pqai_agent.logging_service import logger
from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool


class MessageNotifier(BaseToolkit):
    def __init__(self):
        super().__init__()

    def message_notify_user(self, message: str) -> str:
        """Sends a message to the user.
        Args:
            message (str): The message to send.
        Returns:
            str: A confirmation message.
        """

        logger.info(f"Message to user: {message}")
        return 'success'

    def output_multimodal_message(self, message: Dict[str, str]) -> str:
        """Outputs a multimodal message to the user.
        Args:
            message (Dict[str, str]): The message to output. Message schema:
            {
                "type": "text|image|gif|video|mini_program",
                "content": "message content",
                "title": "only needed if type is video or mini_program",
                "cover_image": "only needed if type is mini_program",
            }
            if message type is image, gif, video or mini_program, the content should be a URL.
        Returns:
            str: A confirmation message.
        """
        logger.info(f"Multimodal message to user: {message}")
        return 'success'


    def get_tools(self):
        return [FunctionTool(self.message_notify_user),
                FunctionTool(self.output_multimodal_message)]