nxy/bot/quotes.py

# -*- coding: utf-8 -*-
import re

from docopt import Dict
from irc3.plugins.command import Commands, command
from irc3.utils import IrcString
from psycopg2 import Error

from . import DatabasePlugin, Bot
from .utils import parse_int


class Quotes(DatabasePlugin):
    def __init__(self, bot: Bot):
        super().__init__(bot)
        self.guard = bot.get_plugin(Commands).guard

    def add_quote(self, mask: IrcString, nick: str, quote: str, channel: IrcString):
        # Parse nick from "<@foobar>" like strings
        nick = re.match(r'<?[~&@%+]?([a-zA-Z0-9_\-^`|\\\[\]{}]+)>?', nick).group(1)

        if not nick:
            self.bot.notice(mask.nick, '[Quotes] Error parsing nick')
        else:
            # Insert quote into database
            self.cur.execute('''
            INSERT INTO
              quotes (nick, item, channel, created_by)
            VALUES
              (%s, %s, %s, %s)
            ''', [nick, quote, channel, mask.nick])

    def delete_quote(self, nick: str, quote: str):
        index, order = parse_int(quote, select=False)

        if index:
            # Delete from database
            self.cur.execute('''
            -- noinspection SqlResolve
            WITH ranked_quotes AS (
              SELECT
                id,
                rank() OVER (PARTITION BY nick ORDER BY id {order})
              FROM
                quotes
              WHERE
                lower(nick) = lower(%s)
            )

            -- noinspection SqlResolve
            DELETE FROM
              quotes
            WHERE
              id = (
                SELECT
                  id
                FROM
                  ranked_quotes
                WHERE
                  rank = %s
              )
            '''.format(order=order), [nick, index])

    @command(options_first=True, quiet=True)
    def q(self, mask: IrcString, target: IrcString, args: Dict):
        """Get, add or delete quotes for an user

        %%q <cmd> <nick> <quote>...
        %%q <nick> [<index>]
        """
        cmd = args.get('<cmd>')
        nick = args['<nick>']
        quote = ' '.join(args.get('<quote>'))

        if cmd and quote:
            try:
                # Anybody can add
                if cmd == 'add':
                    self.add_quote(mask, nick, quote, target)
                # But only admins can delete
                elif cmd == 'del' and self.guard.has_permission(mask, 'admin'):
                    self.delete_quote(nick, quote)

                self.con.commit()
            except Error as ex:
                # Rollback transaction on error
                print(ex)
                self.con.rollback()
        else:
            index = args.get('<index>')
            values = [nick]

            if index:
                index = parse_int(index)
                if not index:
                    return
                index, order = index
                order = 'rank {order}'.format(order=order)
                offset = 'OFFSET %s'
                values.append(index)
            else:
                order = 'random()'
                offset = ''

            # Fetch quote from database
            self.cur.execute('''
            WITH ranked_quotes AS (
              SELECT
                nick,
                item,
                rank() OVER (PARTITION BY nick ORDER BY id),
                count(*) OVER (PARTITION BY nick) AS total
              FROM
                quotes
            )

            SELECT
              *
            FROM
              ranked_quotes
            WHERE
              lower(nick) LIKE lower(%s)
            ORDER BY
              {order}
            LIMIT
              1
            {offset}
            '''.format(order=order, offset=offset), values)
            result = self.cur.fetchone()

            if result:
                return '[{rank}/{total}] <{nick}> {item}'.format(**result)