croccybot/bot/modules/pending-rewrite/study/badges/badge_tracker.py

import datetime
import traceback
import logging
import asyncio
import contextlib

import discord

from meta import client, sharding
from data.conditions import GEQ, THIS_SHARD
from core.data import lions
from utils.lib import strfdur
from settings import GuildSettings

from ..module import module
from .data import new_study_badges, study_badges


guild_locks = {}  # guildid -> Lock


@contextlib.asynccontextmanager
async def guild_lock(guildid):
    """
    Per-guild lock held while the study badges are being updated.
    This should not be used to lock the data modifications, as those are synchronous.
    Primarily for reporting and so that the member information (e.g. roles) stays consistent
    through reading and manipulation.
    """
    # Create the lock if it hasn't been registered already
    if guildid in guild_locks:
        lock = guild_locks[guildid]
    else:
        lock = guild_locks[guildid] = asyncio.Lock()

    await lock.acquire()
    try:
        yield lock
    finally:
        lock.release()


async def update_study_badges(full=False):
    while not client.is_ready():
        await asyncio.sleep(1)

    client.log(
        "Running global study badge update.".format(
        ),
        context="STUDY_BADGE_UPDATE",
        level=logging.DEBUG
    )
    # TODO: Consider db procedure for doing the update and returning rows

    # Retrieve member rows with out of date study badges
    if not full and client.appdata.last_study_badge_scan is not None:
        # TODO: _extra here is a hack to cover for inflexible conditionals
        update_rows = new_study_badges.select_where(
            guildid=THIS_SHARD,
            _timestamp=GEQ(client.appdata.last_study_badge_scan or 0),
            _extra="OR session_start IS NOT NULL AND (guildid >> 22) %% {} = {}".format(
                sharding.shard_count, sharding.shard_number
            )
        )
    else:
        update_rows = new_study_badges.select_where(guildid=THIS_SHARD)

    if not update_rows:
        client.appdata.last_study_badge_scan = datetime.datetime.utcnow()
        return

    # Batch and fire guild updates
    current_guildid = None
    current_guild = None
    guild_buffer = []
    updated_guilds = set()
    for row in update_rows:
        if row['guildid'] != current_guildid:
            if current_guild:
                # Fire guild updater
                asyncio.create_task(_update_guild_badges(current_guild, guild_buffer))
                updated_guilds.add(current_guild.id)

            guild_buffer = []
            current_guildid = row['guildid']
            current_guild = client.get_guild(row['guildid'])

        if current_guild:
            guild_buffer.append(row)

    if current_guild:
        # Fire guild updater
        asyncio.create_task(_update_guild_badges(current_guild, guild_buffer))
        updated_guilds.add(current_guild.id)

    # Update the member study badges in data
    lions.update_many(
        *((row['current_study_badgeid'], row['guildid'], row['userid'])
          for row in update_rows if row['guildid'] in updated_guilds),
        set_keys=('last_study_badgeid',),
        where_keys=('guildid', 'userid'),
        cast_row='(NULL::int, NULL::int, NULL::int)'
    )

    # Update the app scan time
    client.appdata.last_study_badge_scan = datetime.datetime.utcnow()


async def _update_guild_badges(guild, member_rows, notify=True, log=True):
    """
    Notify, update, and log role changes for a single guild.
    Expects a valid `guild` and a list of Rows of `new_study_badges`.
    """
    async with guild_lock(guild.id):
        client.log(
            "Running guild badge update for guild '{guild.name}' (gid:{guild.id}) "
            "with `{count}` rows to update.".format(
                guild=guild,
                count=len(member_rows)
            ),
            context="STUDY_BADGE_UPDATE",
            level=logging.DEBUG,
            post=False
        )

        # Set of study role ids in this guild, usually from cache
        guild_roles = {
            roleid: guild.get_role(roleid)
            for roleid in study_badges.queries.for_guild(guild.id)
        }

        log_lines = []
        flags_used = set()
        tasks = []
        for row in member_rows:
            # Fetch member
            # TODO: Potential verification issue
            member = guild.get_member(row['userid'])

            if member:
                tasks.append(
                    asyncio.create_task(
                        _update_member_roles(row, member, guild_roles, log_lines, flags_used, notify)
                    )
                )

        # Post to the event log, in multiple pages if required
        event_log = GuildSettings(guild.id).event_log.value
        if tasks:
            task_blocks = (tasks[i:i+20] for i in range(0, len(tasks), 20))
            for task_block in task_blocks:
                # Execute the tasks
                await asyncio.gather(*task_block)

                # Post to the log if needed
                if log and event_log:
                    desc = "\n".join(log_lines)
                    embed = discord.Embed(
                        title="Study badge{} earned!".format('s' if len(log_lines) > 1 else ''),
                        description=desc,
                        colour=discord.Colour.orange(),
                        timestamp=datetime.datetime.utcnow()
                    )
                    if flags_used:
                        flag_desc = {
                            '!': "`!` Could not add/remove badge role. **Check permissions!**",
                            '*': "`*` Could not message member.",
                            'x': "`x` Couldn't find role to add/remove!"
                        }
                        flag_lines = '\n'.join(desc for flag, desc in flag_desc.items() if flag in flags_used)
                        embed.add_field(
                            name="Legend",
                            value=flag_lines
                        )
                    try:
                        await event_log.send(embed=embed)
                    except discord.HTTPException:
                        # Nothing we can really do
                        pass

                # Flush the log collection pointers
                log_lines.clear()
                flags_used.clear()

                # Wait so we don't get ratelimited
                await asyncio.sleep(0.5)

    # Debug log completion
    client.log(
        "Completed guild badge update for guild '{guild.name}' (gid:{guild.id})".format(
            guild=guild,
        ),
        context="STUDY_BADGE_UPDATE",
        level=logging.DEBUG,
        post=False
    )


async def _update_member_roles(row, member, guild_roles, log_lines, flags_used, notify):
    guild = member.guild

    # Logging flag chars
    flags = []

    # Add new study role
    # First fetch the roleid using the current_study_badgeid
    new_row = study_badges.fetch(row['current_study_badgeid']) if row['current_study_badgeid'] else None

    # Fetch actual role from the precomputed guild roles
    to_add = guild_roles.get(new_row.roleid, None) if new_row else None
    if to_add:
        # Actually add the role
        try:
            await member.add_roles(
                to_add,
                atomic=True,
                reason="Updating study badge."
            )
        except discord.HTTPException:
            flags.append('!')
    elif new_row:
        flags.append('x')

    # Remove other roles, start by trying the last badge role
    old_row = study_badges.fetch(row['last_study_badgeid']) if row['last_study_badgeid'] else None

    member_roleids = set(role.id for role in member.roles)
    if old_row and old_row.roleid in member_roleids:
        # The last level role exists, try to remove it
        try:
            await member.remove_roles(
                guild_roles.get(old_row.roleid),
                atomic=True
            )
        except discord.HTTPException:
            # Couldn't remove the role
            flags.append('!')
    else:
        # The last level role doesn't exist or the member doesn't have it
        # Remove all leveled roles they have
        current_roles = (
            role for roleid, role in guild_roles.items()
            if roleid in member_roleids and (not to_add or roleid != to_add.id)
        )
        if current_roles:
            try:
                await member.remove_roles(
                    *current_roles,
                    atomic=True,
                    reason="Updating study badge."
                )
            except discord.HTTPException:
                # Couldn't remove one or more of the leveled roles
                flags.append('!')

    # Send notification to member
    # TODO: Config customisation
    if notify and new_row and (old_row is None or new_row.required_time > old_row.required_time):
        req = new_row.required_time
        if req < 3600:
            timestr = "{} minutes".format(int(req // 60))
        elif req == 3600:
            timestr = "1 hour"
        elif req % 3600:
            timestr = "{:.1f} hours".format(req / 3600)
        else:
            timestr = "{} hours".format(int(req // 3600))
        embed = discord.Embed(
            title="New Study Badge!",
            description="Congratulations! You have earned {} for studying **{}**!".format(
                "**{}**".format(to_add.name) if to_add else "a new study badge!",
                timestr
            ),
            timestamp=datetime.datetime.utcnow(),
            colour=discord.Colour.orange()
        ).set_footer(text=guild.name, icon_url=guild.icon_url)
        try:
            await member.send(embed=embed)
        except discord.HTTPException:
            flags.append('*')

    # Add to event log message
    if new_row:
        new_role_str = "earned <@&{}> **({})**".format(new_row.roleid, strfdur(new_row.required_time))
    else:
        new_role_str = "lost their study badge!"
    log_lines.append(
        "<@{}> {} {}".format(
            row['userid'],
            new_role_str,
            "`[{}]`".format(''.join(flags)) if flags else "",
        )
    )
    if flags:
        flags_used.update(flags)


async def study_badge_tracker():
    """
    Runloop for the study badge updater.
    """
    while True:
        try:
            await update_study_badges()
        except Exception:
            # Unknown exception. Catch it so the loop doesn't die.
            client.log(
                "Error while updating study badges! "
                "Exception traceback follows.\n{}".format(
                    traceback.format_exc()
                ),
                context="STUDY_BADGE_TRACKER",
                level=logging.ERROR
            )
        # Long delay since this is primarily needed for external modifications
        # or badge updates while studying
        await asyncio.sleep(60)


async def update_member_studybadge(member):
    """
    Checks and (if required) updates the study badge for a single member.
    """
    update_rows = new_study_badges.select_where(
        guildid=member.guild.id,
        userid=member.id
    )
    if update_rows:
        # Debug log the update
        client.log(
            "Updating study badge for user '{member.name}' (uid:{member.id}) "
            "in guild '{member.guild.name}' (gid:{member.guild.id}).".format(
                member=member
            ),
            context="STUDY_BADGE_UPDATE",
            level=logging.DEBUG
        )

        # Update the data first
        lions.update_where({'last_study_badgeid': update_rows[0]['current_study_badgeid']},
                           guildid=member.guild.id, userid=member.id)

        # Run the update task
        await _update_guild_badges(member.guild, update_rows)


@module.launch_task
async def launch_study_badge_tracker(client):
    asyncio.create_task(study_badge_tracker())