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 is_int, 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
            with self.con.cursor() as cur:
                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
            with self.con.cursor() as cur:
                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 = args.get('<quote>')

        if (cmd == 'add' or cmd == 'del') and quote:
            quote = ' '.join(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:
            query = None
            index = args.get('<index>')
            where = []
            values = []

            # search query support
            if cmd or index and not is_int(index):
                if cmd:
                    query = nick
                    nick = cmd
                else:
                    query = index

                # if last entry in quotes list is a number, use it as the index
                if quote and is_int(quote[-1]):
                    index = quote[-1]
                    quote = quote[:-1]
                # else get random quote
                else:
                    index = None

                if quote:
                    query += ' ' + ' '.join(quote)
            else:
                quote = ' '.join(quote)

            if nick != 'search' or not query:
                where.append('nick ILIKE %s')
                values.append(nick)

            if query:
                where.append('item ILIKE \'%%\' || %s || \'%%\'')
                values.append(query)

            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
            with self.con.cursor() as cur:
                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
                  {where}
                ORDER BY
                  {order}
                LIMIT
                  1
                {offset}
                '''.format(where=' AND '.join(where), order=order, offset=offset), values)
                result = cur.fetchone()

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