client.py

from __future__ import annotations

import asyncio
import json
import pprint
import sys
import traceback
import logging
from pprint import pformat
from typing import List, Any, AsyncGenerator, Dict, Coroutine, TypeVar
from typing import Optional

from aiohttp import web
from slackclient import SlackClient

import hooks
import settings
import slack_util

"""
Objects to wrap slack connections
"""
# Read the API token


api_file = open("apitoken.txt", 'r')
SLACK_API = next(api_file).strip()
api_file.close()


class ClientWrapper(object):
    """
    Essentially the main state object.
    We only ever expect one of these per api token.
    Holds a slack client, and handles messsages.
    """

    def __init__(self, api_token):
        # Init slack
        self.slack = SlackClient(api_token)

        # Hooks go regex -> callback on (slack, msg, match)
        self.hooks: List[hooks.AbsHook] = []

        # Periodicals are just wrappers around an iterable, basically
        self.passives: List[hooks.Passive] = []

        # Cache users and channels
        self.users: Dict[str, slack_util.User] = {}
        self.conversations: Dict[str, slack_util.Conversation] = {}

    # Scheduled/passive events handling
    def add_passive(self, per: hooks.Passive) -> None:
        self.passives.append(per)

    async def run_passives(self) -> None:
        """
        Run all currently added passives
        """
        awaitables = [p.run() for p in self.passives]
        await asyncio.gather(*awaitables)

    # Incoming slack hook handling
    def add_hook(self, hook: hooks.AbsHook) -> None:
        self.hooks.append(hook)

    async def handle_events(self) -> None:
        """
        Asynchronous tasks that eternally reads and responds to messages.
        """
        # Create a queue
        queue = asyncio.Queue()

        # Create a task to put rtm events to the queue
        rtm_task = asyncio.create_task(self.rtm_event_feed(queue))

        # Create a task to put http events to the queue
        http_task = asyncio.create_task(self.http_event_feed(queue))

        # Create a task to handle all other tasks
        async def handle_task_loop():
            async for t3 in self.spool_tasks(queue):
                sys.stdout.flush()
                if settings.SINGLE_THREAD_TASKS:
                    await t3

        # Handle them all
        await asyncio.gather(rtm_task, http_task, handle_task_loop())

    async def rtm_event_feed(self, msg_queue: asyncio.Queue) -> None:
        """
        Async wrapper around the message feed.
        Yields messages awaitably forever.
        """
        # Create the msg feed
        feed = slack_util.message_stream(self.slack)

        # Create a simple callable that gets one message from the feed
        def get_one():
            return next(feed)

        # Continuously yield async threaded tasks that poll the feed
        while True:
            next_event = await asyncio.get_running_loop().run_in_executor(None, get_one)
            await msg_queue.put(next_event)

    async def http_event_feed(self, event_queue: asyncio.Queue) -> None:
        # Create a callback to convert requests to events
        async def interr(request: web.Request):
            if request.can_read_body:
                # Get the payload
                post_params = await request.post()
                payload = json.loads(post_params["payload"])
                logging.info("\nInteraction Event received:")
                logging.debug(pformat(payload))

                # Handle each action separately
                if "actions" in payload:
                    for action in payload["actions"]:
                        # Start building the event
                        ev = slack_util.Event()

                        # Get the user who clicked the button
                        ev.user = slack_util.UserContext(payload["user"]["id"])

                        # Get the message that they clicked
                        ev.message = slack_util.RelatedMessageContext(payload["message"]["ts"], payload["message"]["text"])

                        # Get the channel it was clicked in
                        ev.conversation = slack_util.ConversationContext(payload["channel"]["id"])

                        # Get the message this button/action was attached to
                        ev.interaction = slack_util.InteractionContext(payload["response_url"],
                                                                       payload["trigger_id"],
                                                                       action["block_id"],
                                                                       action["action_id"],
                                                                       action.get("value"))

                        # Put it in the queue
                        await event_queue.put(ev)

                # Respond that everything is fine
                return web.Response(status=200)
            else:
                logging.error("\nMalformed event received.")
                # If we can't read it, get mad
                return web.Response(status=400)

        # Create the server
        app = web.Application()
        app.add_routes([web.post('/bothttpcallback', interr)])

        # Asynchronously serve that boy up
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, port=31019)
        await site.start()
        logging.info("Server up")
        # while True:
        #     await asyncio.sleep(30)

    async def spool_tasks(self, event_queue: asyncio.Queue) -> AsyncGenerator[asyncio.Task, Any]:
        """
        Read in from async event feed, and spool them out as async tasks
        """
        while True:
            event: slack_util.Event = await event_queue.get()
            # Find which hook, if any, satisfies
            for hook in list(self.hooks):  # Note that we do list(self.hooks) to avoid edit-while-iterating issues
                # Try invoking each
                try:
                    # Try to make a coroutine handling the message
                    coro = hook.try_apply(event)

                    # If we get a coro back, then task it up and set consumption appropriately
                    if coro is not None:
                        logging.debug("Spawned task. Now {} running total.".format(len(asyncio.all_tasks())))
                        yield asyncio.create_task(_exception_printing_task(coro))
                        if hook.consumes:
                            break

                except hooks.HookDeath:
                    # If a hook wants to die, let it.
                    self.hooks.remove(hook)

    # Data getting/sending

    def get_conversation(self, conversation_id: str) -> Optional[slack_util.Conversation]:
        return self.conversations.get(conversation_id)

    def get_conversation_by_name(self, conversation_identifier: str) -> Optional[slack_util.Conversation]:
        # If looking for a direct message, first lookup user, then fetch
        if conversation_identifier[0] == "@":
            # user_name = conversation_identifier

            # Find the user by their name
            raise NotImplementedError("There wasn't a clear use case for this yet, so we've opted to just not use it")

        # If looking for a channel, just lookup normally
        elif conversation_identifier[0] == "#":
            channel_name = conversation_identifier

            # Find the channel in the dict
            for channel in self.conversations.values():
                if channel.name == channel_name:
                    return channel

        # If it doesn't fit the above, we don't know how to process
        else:
            raise ValueError("Please give either an #channel-name or @user-name")

        # If we haven't returned already, give up and return None
        return None

    def get_user(self, user_id: str) -> Optional[slack_util.User]:
        return self.users.get(user_id)

    def get_user_by_name(self, user_name: str) -> Optional[slack_util.User]:
        raise NotImplementedError()

    def api_call(self, api_method, **kwargs):
        return self.slack.api_call(api_method, **kwargs)

    # Simpler wrappers around message sending/replying

    def reply(self, event: slack_util.Event, text: str, in_thread: bool = True) -> dict:
        """
        Replies to a message.
        Message must have a channel and message context.
        Returns the JSON response.
        """
        # Ensure we're actually replying to a valid message
        assert (event.conversation and event.message) is not None

        # Send in a thread by default
        if in_thread:
            # Figure otu what thread to send it to
            thread = event.message.ts
            if event.thread:
                thread = event.thread.thread_ts
            return self.send_message(text, event.conversation.conversation_id, thread=thread)
        else:
            return self.send_message(text, event.conversation.conversation_id)

    def _send_core(self, api_method: str, text: str, channel_id: str, thread: str, broadcast: bool,
                   blocks: Optional[List[dict]]) -> dict:
        """
        Copy of the internal send message function of slack, with some helpful options.
        Returns the JSON response.
        """
        # Check that text exists if there are no blocks
        if blocks is None and text is None:
            raise ValueError("Must provide blocks or texts or both.")
        elif text is None:
            text = "_Block message. Open slack client to view_"

        # Begin constructing kwargs with fields that _must_ exist
        kwargs = {"channel": channel_id, "text": text, "as_user": True, "parse": True}

        # Deduce thread stuff
        if thread:
            kwargs["thread_ts"] = thread
            if broadcast:
                kwargs["reply_broadcast"] = True
        elif broadcast:
            raise ValueError("Can't broadcast a non-threaded message. Try again.")

        # Set blocks iff provided.
        if blocks is not None:
            kwargs["blocks"] = blocks

        result = self.api_call(api_method, **kwargs)

        logging.info("Tried to send message \"{}\". Got response:\n {}".format(text, pprint.pformat(result)))
        return result

    def send_message(self,
                     text: Optional[str],
                     channel_id: str,
                     thread: str = None,
                     broadcast: bool = False,
                     blocks: Optional[List[dict]] = None) -> dict:
        """
        Wraps _send_core for normal messages
        """
        return self._send_core("chat.postMessage", text, channel_id, thread, broadcast, blocks)

    def send_ephemeral(self,
                       text: Optional[str],
                       channel_id: str,
                       thread: str = None,
                       blocks: Optional[List[dict]] = None) -> dict:
        """
        Wraps _send_core for ephemeral messages
        """
        return self._send_core("chat.postEphemeral", text, channel_id, thread, False, blocks)

    def edit_message(self, text: Optional[str], channel_id: str, message_ts: str, blocks: Optional[List[dict]] = None):
        """
        Edits a message.
        """
        # Check that text exists if there are no blocks
        if blocks is None and text is None:
            raise ValueError("Must provide blocks or texts or both.")
        elif text is None:
            text = "_Block message. Open slack client to view_"

        # Set the default args
        kwargs = {"ts": message_ts, "channel": channel_id, "text": text, "as_user": True}

        # Set blocks if provided
        if blocks is not None:
            kwargs["blocks"] = blocks

        return self.api_call("chat.update", **kwargs)

    # Update slack data

    def update_channels(self):
        """
        Queries the slack API for all current channels
        """
        # Necessary because of pagination
        cursor = None

        # Make a new dict to use
        new_dict = {}

        # Iterate over results
        while True:
            # Set args depending on if a cursor exists
            args = {"limit": 1000, "types": "public_channel,private_channel,mpim,im"}
            if cursor:
                args["cursor"] = cursor

            channel_dicts = self.api_call("conversations.list", **args)

            # If the response is good, put its results to the dict
            if channel_dicts["ok"]:
                for channel_dict in channel_dicts["channels"]:
                    if channel_dict["is_im"]:
                        new_channel = slack_util.DirectMessage(id=channel_dict["id"],
                                                               user_id="@" + channel_dict["user"])
                    else:
                        new_channel = slack_util.Channel(id=channel_dict["id"],
                                                         name="#" + channel_dict["name"])
                    new_dict[new_channel.id] = new_channel

                # Fetch the cursor
                cursor = channel_dicts.get("response_metadata").get("next_cursor")

                # If cursor is blank, we're done new channels, just give it up
                if cursor == "":
                    break

            else:
                logging.warning("Failed to retrieve channels. Message: {}".format(channel_dicts))
                break
        self.conversations = new_dict

    def update_users(self):
        """
        Queries the slack API for all current users
        """
        # Necessary because of pagination
        cursor = None

        while True:
            # Set args depending on if a cursor exists
            args = {"limit": 1000}
            if cursor:
                args["cursor"] = cursor

            user_dicts = self.api_call("users.list", **args)

            # Make a new dict to use
            new_dict = {}

            # If the response is good:
            if user_dicts["ok"]:
                for user_dict in user_dicts["members"]:
                    new_user = slack_util.User(id=user_dict.get("id"),
                                               name=user_dict.get("name"),
                                               real_name=user_dict.get("real_name"),
                                               email=user_dict.get("profile").get("email"))
                    new_dict[new_user.id] = new_user

                # Fetch the cursor
                cursor = user_dicts.get("response_metadata").get("next_cursor")

                # If cursor is blank, we're done new channels, just give it up
                if cursor == "":
                    break

            else:
                logging.warning("Warning: failed to retrieve users")
                break
        self.users = new_dict


# Create a single instance of the client wrapper
_singleton = ClientWrapper(SLACK_API)


def get_slack() -> ClientWrapper:
    return _singleton


"""
Miscellania
"""

A = TypeVar("A")
B = TypeVar("B")
C = TypeVar("C")


# Prints exceptions instead of silently dropping them in async tasks
async def _exception_printing_task(c: Coroutine[A, B, C]) -> Coroutine[A, B, C]:
    # Print exceptions as they pass through
    try:
        return await c
    except Exception:
        output = traceback.format_exc()
        print(output)
        get_slack().send_message(output, "#botzone")
        raise