scheduler.py

import discord
from discord.ext import commands
from cogs.utils import checks
from cogs.utils.dataIO import dataIO
from cogs.utils.chat_formatting import box, error, warning

import logging
import os
import re
import asyncio
import time
from datetime import datetime, timezone, timedelta
from random import randint
from math import ceil
from collections import defaultdict

__version__ = '2.0.5'

log = logging.getLogger("red.scheduler")
log.setLevel(logging.INFO)

PATH = 'data/scheduler/'
JSON = PATH + 'events.json'

UNIT_TABLE = (
    (('weeks', 'wks', 'w'), 60 * 60 * 24 * 7),
    (('days', 'dys', 'd'), 60 * 60 * 24),
    (('hours', 'hrs', 'h'), 60 * 60),
    (('minutes', 'mins', 'm'), 60),
    (('seconds', 'secs', 's'), 1),
)


class Event:
    __slots__ = ['name', 'channel', 'server', 'author', 'command', 'timedelta', 'repeat', 'starttime']

    def __init__(self, **data):
        self.name = data['name']
        self.channel = data['channel']
        self.server = data['server']
        self.author = data['author']
        self.command = data['command']
        self.timedelta = data['timedelta']
        self.repeat = data['repeat']
        self.starttime = data.get('starttime')

    @staticmethod
    def _key(obj):
        return (obj.timedelta, obj.name, obj.starttime, obj.channel)

    def __hash__(self):
        return hash(self._key(self))

    def __eq__(self, other):
        if not isinstance(other, Event):
            return False

        for k in self.__slots__:
            if getattr(self, k) != getattr(other, k):
                return False

        return True

    def __lt__(self, other):
        if not isinstance(other, Event):
            raise TypeError('only Event can be compared')

        return Event._key(self) < Event._key(other)


class BadTimeExpr(ValueError):
    pass


def _find_unit(unit):
    for names, length in UNIT_TABLE:
        if any(n.startswith(unit) for n in names):
            return names, length
    raise BadTimeExpr("Invalid unit: %s" % unit)


def _parse_time(time):
    time = time.lower()
    if not time.isdigit():
        time = re.split(r'\s*([\d.]+\s*[^\d\s,;]*)(?:[,;\s]|and)*', time)
        time = sum(map(_timespec_sec, filter(None, time)))
    return int(time)


def _timespec_sec(expr):
    atoms = re.split(r'([\d.]+)\s*([^\d\s]*)', expr)
    atoms = list(filter(None, atoms))

    if len(atoms) > 2:  # This shouldn't ever happen
        raise BadTimeExpr("invalid expression: '%s'" % expr)
    elif len(atoms) == 2:
        names, length = _find_unit(atoms[1])
        if atoms[0].count('.') > 1 or \
                not atoms[0].replace('.', '').isdigit():
            raise BadTimeExpr("Not a number: '%s'" % atoms[0])
    else:
        names, length = _find_unit('seconds')

    try:
        return float(atoms[0]) * length
    except ValueError:
        raise BadTimeExpr("invalid value: '%s'" % atoms[0])


def _generate_timespec(sec, short=False, micro=False):
    timespec = []

    for names, length in UNIT_TABLE:
        n, sec = divmod(sec, length)

        if n:
            if micro:
                s = '%d%s' % (n, names[2])
            elif short:
                s = '%d%s' % (n, names[1])
            else:
                s = '%d %s' % (n, names[0])
            if n <= 1:
                s = s.rstrip('s')
            timespec.append(s)

    if len(timespec) > 1:
        if micro:
            return ''.join(timespec)

        segments = timespec[:-1], timespec[-1:]
        return ' and '.join(', '.join(x) for x in segments)

    return timespec[0]


def _convert_iso8601(input_string):
    tsre = r"[:]|([-](?!((\d{2}[:]\d{2})|(\d{4}))$))"
    ts = re.sub(tsre, '', input_string)

    if ts.endswith(('z', 'Z')):
        ts = ts[:-1] + '+0000'

    if '.' in ts:
        fmt = "%Y%m%dT%H%M%S.%f%z"
    else:
        fmt = "%Y%m%dT%H%M%S%z"

    return datetime.strptime(ts, fmt)


class Scheduler:
    """Schedules commands to run every so often.

    Times are formed as follows: 1s, 2m, 3h, 5d, 1w
    """

    def __init__(self, bot):
        self.bot = bot
        self.events = dataIO.load_json(JSON)
        self.queue = asyncio.PriorityQueue(loop=self.bot.loop)
        self.queue_lock = asyncio.Lock()
        self.pending = {}
        self.pending_by_event = defaultdict(lambda: list())
        self._load_events()
        self.task = bot.loop.create_task(self.queue_manager())

    def __unload(self):
        self.task.cancel()

    def save_events(self):
        dataIO.save_json(JSON, self.events)

    def _load_events(self):
        # for entry in the self.events make an Event
        for server in self.events:
            for name, event in self.events[server].items():
                e = Event(server=server, **event)
                self.bot.loop.create_task(self._put_event(e))

    async def _put_event(self, event, fut=None):
        # compute if absolute time isn't provided
        if fut is None:
            if event.repeat:
                diff = max(time.time() - event.starttime, 0)
                fut = ((ceil(diff / event.timedelta) * event.timedelta) + event.starttime)
            else:
                fut = event.starttime + event.timedelta

        await self.queue.put((fut, event))

        log.debug('Added "{}" to the scheduler queue at {}'.format(event.name, fut))

    async def _add_event(self, name, command, dest_server, dest_channel, author, timedelta, repeat=False, start=None):
        if isinstance(dest_server, discord.Server):
            dest_server = dest_server.id

        if isinstance(dest_channel, discord.Channel):
            dest_channel = dest_channel.id

        if isinstance(author, discord.User):
            author = author.id

        if dest_server not in self.events:
            self.events[dest_server] = {}

        if isinstance(start, datetime):
            start = start.timestamp()

        event_dict = {
            'name'      : name,
            'channel'   : dest_channel,
            'author'    : author,
            'command'   : command,
            'timedelta' : timedelta,
            'repeat'    : repeat,
            'starttime' : start or int(time.time())
        }

        log.debug('event dict:\n\t{}'.format(event_dict))

        self.events[dest_server][name] = event_dict
        e = Event(server=dest_server, **event_dict)
        await self._put_event(e)

        self.save_events()

    async def _remove_event(self, name, server):
        events = []
        removed = None

        async with self.queue_lock:

            while not self.queue.empty():
                time, event = self.queue.get_nowait()

                if name == event.name and server.id == event.server:
                    removed = (time, event)
                    break
                else:
                    events.append((time, event))

            for event in events:
                self.queue.put_nowait(event)

        if not removed:
            return None

        time, event = removed

        if self.pending_by_event[event]:
            for ts in self.pending_by_event[event]:
                k = (ts, event)
                self.pending[k].cancel()
                del self.pending[k]

            del self.pending_by_event[event]

        return removed

    @commands.group(no_pm=True, pass_context=True)
    async def scheduler(self, ctx):
        if ctx.invoked_subcommand is None:
            await self.bot.send_cmd_help(ctx)
            return

    @scheduler.command(pass_context=True, name="add")
    async def _scheduler_add(self, ctx, time_interval, *, command):
        """Add a command to run in [time_interval].

        Intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        name = ctx.message.author.id + '-' + command.lower()
        await self._add_centralized(ctx, name, time_interval, command, False)

    @scheduler.command(pass_context=True, name="add_timelast")
    async def _scheduler_add_timelast(self, ctx, command, *, time_interval):
        """Add a command to run in [time_interval].

        Intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        name = ctx.message.author.id + '-' + command.lower()
        await self._add_centralized(ctx, name, time_interval, command, False)

    @scheduler.command(pass_context=True, name="add_twostage")
    async def _scheduler_add_twostage(self, ctx, command1, time_interval, *, command2):
        """Add a command to run now, and another in [time_interval].

        Intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        name1 = ctx.message.author.id + '-' + command1.lower()
        name2 = ctx.message.author.id + '-' + command2.lower()
        self.dispatch_fake(ctx.message.channel, ctx.message.author.id, name1, command1)
        await self._add_centralized(ctx, name2, time_interval, command2, False)

    @scheduler.command(pass_context=True, name="add_twostage_timelast")
    async def _scheduler_add_twostage_timelast(self, ctx, command1, command2, *, time_interval):
        """Add a command to run now, and another in [time_interval].

        Intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        name1 = ctx.message.author.id + '-' + command1.lower()
        name2 = ctx.message.author.id + '-' + command2.lower()
        self.dispatch_fake(ctx.message.channel, ctx.message.author.id, name1,
                           command1)
        await self._add_centralized(ctx, name2, time_interval, command2, False)

    @scheduler.command(pass_context=True, name="repeat")
    @checks.mod_or_permissions(manage_messages=True)
    async def _scheduler_repeat(self, ctx, name, time_interval, *, command):
        """Add a command to run every [time_interval].

        Intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        await self._add_centralized(ctx, name, time_interval, command, True)

    @scheduler.command(pass_context=True, name="repeat_from")
    @checks.mod_or_permissions(manage_messages=True)
    async def _scheduler_repeat_from(self, ctx, name, start, interval, *, command):
        """Add a command to run every [interval] starting at [start].

        Start time must be an ISO8601 timestamp, unix timestamp or 'now'.

        An ISO8601 timestamp looks like this: 2017-12-11T01:15:03.449371-0500.

        The interval is any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        await self._add_centralized(ctx, name, interval, command, True, start)

    @scheduler.command(pass_context=True, name="repeat_in")
    @checks.mod_or_permissions(manage_messages=True)
    async def _scheduler_repeat_in(self, ctx, name, start_in, interval, *, command):
        """Add a command to run every [interval] starting in [start_in].

        Both intervals are any combination of units: 1s, 2m, 3h, 5d, 1w.
        """
        await self._add_centralized(ctx, name, interval, command, True, start_in=start_in)

    async def _add_centralized(self, ctx, name, interval, command, repeat, start=None, start_in=None):
        channel = ctx.message.channel
        server = ctx.message.server
        author = ctx.message.author
        name = name.lower()
        now = datetime.now(tz=timezone.utc)

        min_time = 30 if repeat else 5

        prefix = await self.get_prefix(ctx.message)
        command = command.lstrip(prefix)

        try:
            interval = _parse_time(interval)
        except BadTimeExpr as e:
            await self.bot.say(error(e.args[0]))
            return

        if start_in:
            try:
                start_in = _parse_time(start_in)
            except BadTimeExpr as e:
                await self.bot.say(error(e.args[0]))
                return
            start = now + timedelta(seconds=start_in)
        elif start:
            try:
                start = self._get_start(start, now)
            except ValueError:
                await self.bot.say(error('Invalid timestamp format!'))
                return
        else:
            start = now

        if interval < min_time:
            await self.bot.say("I'm sorry, {}, I can't let you do that. "
                               "Your time interval is waaaay too short and "
                               "I'll likely get rate limited. Try going above"
                               " {} seconds.".format(author.name, min_time))
            return
        elif name in self.events.get(server.id, {}):
            if repeat:
                msg = warning("An event with that name already exists!")
            else:
                msg = warning("That command is already scheduled to run!")

            await self.bot.say(msg)
            return

        if repeat:
            logmsg = 'add {} "{}" to {} on {} every {}s starting {}'
            msg = '"{}" will run `{}` every {}, starting at {} ({}).'
        else:
            logmsg = 'add {} "{}" to {} on {} in {}s'
            msg = 'I will run `{1}` in {2}.'

        log.info(logmsg.format(name, command, channel.name, server.name, interval, start.timestamp()))

        await self._add_event(name, command, server, channel, author, interval, repeat, start)

        timeexpr = _generate_timespec(interval)

        delta = self._format_start(start, now)

        msg = msg.format(name, command, timeexpr, start, delta)
        await self.bot.say(msg)

    @scheduler.command(pass_context=True, name="remove")
    async def _scheduler_remove(self, ctx, *, name):
        """Removes a scheduled repeating command by name."""
        server = ctx.message.server
        name = name.lower()

        if not self.events.get(server.id):
            await self.bot.say('No events are scheduled for this server.')
            return

        del self.events[server.id][name]
        await self._remove_event(name, server)
        self.save_events()
        await self.bot.say('"{}" has successfully been removed.'.format(name))

    @scheduler.command(pass_context=True, name="cancel")
    async def _scheduler_cancel(self, ctx, *, command):
        """Cancels a scheduled oneshot (non-repeating) command."""
        server = ctx.message.server
        fname = ctx.message.author.id + '-' + command.lower()

        event = self.events.get(server.id, {}).pop(fname, None)
        cancelled = await self._remove_event(fname, server)

        if event:
            self.save_events()

        if event or cancelled:
            await self.bot.say('"{}" has been successfully cancelled.'.format(command))
        else:
            await self.bot.say('Cannot find a scheduled run of "{}".'.format(command))

    @scheduler.command(pass_context=True, name="list")
    @checks.mod_or_permissions(manage_messages=True)
    async def _scheduler_list(self, ctx):
        """Lists all repeated commands
        """
        server = ctx.message.server

        if not self.events.get(server.id):
            await self.bot.say('No events scheduled for this server.')
            return

        mess = "Names:\n\t"
        mess += "\n\t".join(sorted(self.events[server.id].keys()))
        await self.bot.say(box(mess))

    def dispatch_fake(self, channel, author_id, name, command):
        prefix = self.bot.settings.get_prefixes(channel.server)[0]

        data = {}
        timestamp = time.strftime("%Y-%m-%dT%H:%M:%S%z", time.gmtime())
        data['timestamp'] = timestamp
        data['id'] = randint(10**(17), (10**18) - 1)
        data['content'] = prefix + command
        data['channel'] = channel
        data['author'] = {'id': author_id}
        data['nonce'] = randint(-2**32, (2**32) - 1)
        data['channel_id'] = channel.id
        data['reactions'] = []
        fake_message = discord.Message(**data)
        log.info("Running '{0}' in {1.server}/#{1}".format(name, channel))
        self.bot.dispatch('message', fake_message)

    def run_coro(self, event, schedtime):
        channel = self.bot.get_channel(event.channel)

        if channel is None:
            log.debug("Channel no longer found, not running scheduled event.")
        else:
            self.dispatch_fake(channel, event.author, event.name, event.command)

        t = (schedtime, event)

        if t in self.pending:
            del self.pending[t]

        self.pending_by_event[event].remove(schedtime)

    async def process_queue_event(self):
        if self.queue.empty():
            return False

        now = int(time.time())
        next_time, next_event = await self.queue.get()

        diff = max(next_time - now, 0)

        if diff < 10:
            log.debug('scheduling call of "{}" in {}s'.format(next_event.name, diff))

            fut = self.bot.loop.call_later(diff, self.run_coro, next_event, next_time)
            self.pending[(next_time, next_event)] = fut
            self.pending_by_event[next_event].append(next_time)

            if next_event.repeat:
                # schedule next run
                await self._put_event(next_event, next_time + next_event.timedelta)
            else:
                del self.events[next_event.server][next_event.name]
                self.save_events()

            return True
        else:
            log.debug('Will run {} "{}" in {}s'.format(next_event.name, next_event.command, diff))
            await self._put_event(next_event, next_time)
            return False

    async def get_prefix(self, msg, content=None):
        prefixes = self.bot.command_prefix
        if callable(prefixes):
            prefixes = prefixes(self.bot, msg)

            if asyncio.iscoroutine(prefixes):
                prefixes = await prefixes

        content = content or msg.content

        for p in prefixes:
            if content.startswith(p):
                return p

        return None

    async def queue_manager(self):
        try:
            await self.bot.wait_until_ready()

            while self == self.bot.get_cog('Scheduler'):
                while True:
                    async with self.queue_lock:
                        if not await self.process_queue_event():
                            break

                await asyncio.sleep(5)

        except asyncio.CancelledError:
            pass
        finally:
            log.debug('manager dying')

            while not self.queue.empty():
                await self.queue.get()

            for fut in self.pending.values():
                fut.cancel()

    def _get_start(self, start, now):
        if start.lower() == 'now' or start is None:
            return now
        elif start.count('.') == 1 and start.replace('.', '').isdigit():
            ts = datetime.utcfromtimestamp(float(start))
            return ts.replace(tzinfo=timezone.utc)
        elif start:
            return _convert_iso8601(start)

    def _format_start(self, start, now):
        if start == now:
            return 'now'
        else:
            diff_seconds = (now - start).total_seconds()
            delta = _generate_timespec(abs(diff_seconds))

            if diff_seconds < 0:
                return 'in ' + delta
            else:
                return delta + ' ago'


def check_folder():
    if not os.path.exists('data/scheduler'):
        os.mkdir('data/scheduler')


def check_files():
    if not dataIO.is_valid_json(JSON):
        print('Creating empty %s' % JSON)
        dataIO.save_json(JSON, {})


def setup(bot):
    check_folder()
    check_files()
    bot.add_cog(Scheduler(bot))