Add reaction handling

3 files changed, 59 insertions, 20 deletions

diff --git a/bot.py b/bot.py
index f29af92..da7bb34 100644
--- a/bot.py
+++ b/bot.py
@@ -7,6 +7,7 @@ import config
 
 import discord
 from discord.ext import commands
+from random import randint
 
 
 class Cupido(commands.Bot):
@@ -16,6 +17,8 @@ class Cupido(commands.Bot):
         self.lobby_channel = None
         self.pair_channels = []
         self.task = None
+        self.reaction_map = {}
+        self.vote_map = {}
 
     async def await_coroutine(self, co):
         if self.task is not None:
@@ -79,17 +82,28 @@ class Cupido(commands.Bot):
             return False
 
     async def send_panel(self, ctx, user, members):
-        channel = user.dm_channel or await user.create_dm()
+        try:
+            channel = user.dm_channel or await user.create_dm()
+        except discord.errors.HTTPException:
+            print(f'warning: could not send pm to {user}')
+            return
         text = config.PANEL_TEXT
-        n = 0
+        reactions = []
+        emojis = config.EMOJI_POOL[:]
+        reaction_map = {}
         for member in members:
             if member == user:
                 continue
-            name = discord.utils.escape_markdown(discord.utils.escape_mentions(member.nick), ignore_links=False)
-            text += f'\n * {config.EMOJI_POOL[n]} {name}'
-            n += 1
+            name = discord.utils.escape_markdown(member.nick or member.display_name, ignore_links=False)
+            emoji = emojis.pop(randint(0, len(emojis) - 1))
+            text += f'\n * {emoji} {name}'
+            reactions.append(emoji)
+            reaction_map[emoji] = member
+        self.reaction_map[user.id] = reaction_map
         embed = discord.Embed(title=config.PANEL_TITLE, type="rich", description=text, colour=discord.Colour.purple())
-        await channel.send(embed=embed)
+        msg = await channel.send(embed=embed)
+        await await_n(map(msg.add_reaction, reactions))
+        return msg
 
     async def destroy_pair_channels(self, ctx, meta_channel):
         await await_n(map(self.try_delete_channel, self.get_pair_channels_no_cache(ctx, meta_channel)))
@@ -110,6 +124,30 @@ class Cupido(commands.Bot):
             return None
         return meta_channel, lobby_channel
 
+    async def on_reaction_add(self, reaction, user):
+        return await self.toggle_vote(reaction, user)
+
+    async def on_reaction_remove(self, reaction, user):
+        return await self.toggle_vote(reaction, user)
+
+    async def toggle_vote(self, reaction, user):
+        if self.vote_map.get(user.id) is None:
+            if user.id in self.reaction_map:
+                self.vote_map[user.id] = self.reaction_map[user.id][str(reaction.emoji)]
+        else:
+            self.vote_map[user.id] = None
+        await self.update_message_panel(reaction.message, self.vote_map.get(user.id), reaction.emoji)
+
+    async def update_message_panel(self, message, vote, emoji):
+        embeds = message.embeds
+        if not embeds: return None
+        embed = embeds[0]
+        current, after = '*', '☑️'
+        if vote is None:
+            current, after = after, current
+        embed.description = embed.description.replace(f'\n {current} {emoji} ', f'\n {after} {emoji} ')
+        return await message.edit(embed=embed)
+
 
 def main():
     token = getenv(config.TOKEN_ENV_VAR)
diff --git a/commands.py b/commands.py
index 2739445..b974435 100644
--- a/commands.py
+++ b/commands.py
@@ -4,6 +4,7 @@ import random
 
 import config
 
+import discord
 from discord.ext import commands
 
 
@@ -57,34 +58,34 @@ async def shuffle(ctx: commands.Context, channel_size=2):
     if not channels:
         return False
     meta_channel, lobby_channel = channels
-    members = lobby_channel.members[:]
+    members = await await_n(map(ctx.guild.fetch_member, lobby_channel.voice_states.keys()))
     if not members:
         await answer(ctx, 'error: nobody wants to get shuffeled :(')
         return False
     if len(members) < channel_size:
         await answer(ctx, f'error: you are too few people ({len(members)}). Group size is {channel_size}')
         return False
-    # round up the quotient between the member count and the channel size
-    channel_count = (len(members) + channel_size - 1) // channel_size
+    channel_count = len(members) // channel_size
+    excess_slots = list(range(channel_count))
+    random.shuffle(excess_slots)
+    excess_slots = excess_slots[:len(members) - channel_count * channel_size]
     ctx.bot.pair_channels = await ctx.bot.create_pair_channels(ctx, meta_channel, channel_count)
     slots = []
     for i, _ in enumerate(ctx.bot.pair_channels):
         slots.extend([i] * channel_size)
+    slots.extend(excess_slots)
     random.shuffle(slots)
     futures = []
-    group_members = [[] for _ in ctx.bot.pair_channels]
-    slot_members = []
-    for slot in slots:
-        member = members.pop()
+    groups = [[] for _ in range(channel_count)]
+    member_slot_map = []
+    for (member, slot) in zip(members, slots):
         channel = ctx.bot.pair_channels[slot]
         futures.append(member.move_to(channel))
-        group_members[slot].append(member)
-        slot_members.append(member)
-    for member, slot in zip(slot_members, slots):
-        group = group_members[slot]
+        groups[slot].append(member)
+        member_slot_map.append(slot)
+    for i, (member, slot) in enumerate(zip(members, member_slot_map)):
+        group = groups[slot]
         await ctx.bot.send_panel(ctx, member, group)
-    if members:
-        await answer(ctx, 'warning: not all members got shuffeled')
     await await_n(futures)
     return True
 
diff --git a/config.py b/config.py
index f59cb71..d7ea555 100644
--- a/config.py
+++ b/config.py
@@ -19,7 +19,7 @@ DEFAULT_LOOP_TIME = 120
 # time that cycles, that "loop" passes
 DEFAULT_LOOP_COUNT = 3
 
-EMOJI_POOL = [':mouse:', ':broccoli:', ':archery:', ':tent:', ':coffee:']
+EMOJI_POOL = ['👾', '🤖', '👻', '🦧', '🍴', '🎮', '🎷', '😈', '🦄', '🐮', '🌻', '🐘', '🍕', '🦉']
 PANEL_TITLE = f'{NAME.title()} control panel'
 PANEL_TEXT = '''
 You have the choice!