Fix reactions

1 files changed, 37 insertions, 25 deletions

diff --git a/bot.py b/bot.py
index da7bb34..ad2f95b 100644
--- a/bot.py
+++ b/bot.py
@@ -7,7 +7,7 @@ import config
 
 import discord
 from discord.ext import commands
-from random import randint
+from random import sample
 
 
 class Cupido(commands.Bot):
@@ -81,28 +81,33 @@ class Cupido(commands.Bot):
         except discord.errors.NotFound:
             return False
 
+    @staticmethod
+    def generate_panel_text(user, members, emojis, reaction_map=None, selected=None):
+        text = config.PANEL_TEXT
+        emojis_chosen = []
+        for member, emoji in zip(members, emojis):
+            if member.id == user.id:
+                continue
+            name = discord.utils.escape_markdown(member.nick or member.display_name, ignore_links=False)
+            text += f'\n {"☑️" if member.id == selected else "*"} {emoji} {name}'
+            if reaction_map is not None:
+                reaction_map[emoji] = member
+            emojis_chosen.append(emoji)
+        return text, emojis_chosen
+
     async def send_panel(self, ctx, user, members):
         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
-        reactions = []
-        emojis = config.EMOJI_POOL[:]
         reaction_map = {}
-        for member in members:
-            if member == user:
-                continue
-            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
+        emojis = sample(config.EMOJI_POOL, k=len(members))
+        text, emojis = self.generate_panel_text(user, members, emojis, reaction_map)
         self.reaction_map[user.id] = reaction_map
         embed = discord.Embed(title=config.PANEL_TITLE, type="rich", description=text, colour=discord.Colour.purple())
         msg = await channel.send(embed=embed)
-        await await_n(map(msg.add_reaction, reactions))
+        await await_n(map(msg.add_reaction, emojis))
         return msg
 
     async def destroy_pair_channels(self, ctx, meta_channel):
@@ -125,28 +130,35 @@ class Cupido(commands.Bot):
         return meta_channel, lobby_channel
 
     async def on_reaction_add(self, reaction, user):
+        """Overrides event handler in Bot"""
         return await self.toggle_vote(reaction, user)
 
     async def on_reaction_remove(self, reaction, user):
+        """Overrides event handler in Bot"""
         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:
+        if reaction.message.author.id == user.id:
+            return
+        if (user.id not in self.reaction_map) or (str(reaction.emoji) not in self.reaction_map[user.id]):
+            return
+        chosen = self.reaction_map[user.id][str(reaction.emoji)]
+        if self.vote_map.get(user.id) == chosen:
             self.vote_map[user.id] = None
-        await self.update_message_panel(reaction.message, self.vote_map.get(user.id), reaction.emoji)
+        else:
+            self.vote_map[user.id] = self.reaction_map[user.id][str(reaction.emoji)]
+        await self.update_message_panel(reaction.message, self.vote_map.get(user.id), user)
 
-    async def update_message_panel(self, message, vote, emoji):
+    async def update_message_panel(self, message, vote, user):
         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)
+        vote = None if vote is None else vote.id
+        embeds[0].description = self.generate_panel_text(
+                user,
+                list(self.reaction_map[user.id].values()),
+                list(self.reaction_map[user.id].keys()),
+                selected=vote)[0]
+        return await message.edit(embed=embeds[0])
 
 
 def main():