Add shuffle command

2 files changed, 113 insertions, 10 deletions

diff --git a/bot.py b/bot.py
index f97f218..88153cc 100644
--- a/bot.py
+++ b/bot.py
@@ -2,6 +2,15 @@ import asyncio
 import discord
 import config
 from command_utils import command, CommandClient
+from random import shuffle, choice
+from time import sleep
+
+
+async def await_n(lst):
+    lst = list(asyncio.create_task(task) for task in lst)
+    if lst:
+        done, _ = await asyncio.wait(lst)
+        return list(task.result() for task in done)
 
 
 class Client(CommandClient):
@@ -9,6 +18,7 @@ class Client(CommandClient):
         super().__init__(*args, **kwargs)
         self.meta_channel = None
         self.lobby_channel = None
+        self.pair_channels = []
 
     async def on_ready(self):
         print(f'the bot {config.NAME} is logged in as "{self.user}"')
@@ -37,14 +47,32 @@ class Client(CommandClient):
 
     def get_meta_channel(self, ctx):
         return (self.meta_channel
-            or discord.utils.get(ctx.guild.categories,
-                                 name=config.CATEGORY_CHANNEL_NAME))
+                or discord.utils.get(ctx.guild.categories,
+                                     name=config.CATEGORY_CHANNEL_NAME))
 
-    def get_lobby_channel(self, ctx):
+    def get_lobby_channel(self, ctx, meta_channel):
         return (self.lobby_channel
-            or discord.utils.get(ctx.guild.voice_channels,
-                                 name=config.LOBBY_CHANNEL_NAME,
-                                 category=self.meta_channel))
+                or discord.utils.get(ctx.guild.voice_channels,
+                                     name=config.LOBBY_CHANNEL_NAME,
+                                     category=meta_channel))
+
+    def get_pair_channels(self, ctx, meta_channel):
+        return (self.pair_channels
+                or sorted((channel for channel in ctx.guild.voice_channels
+                           if channel.category == meta_channel
+                           and channel.name.isdigit()),
+                          key=lambda c: c.name))
+
+    async def destroy_pair_channels(self, ctx, meta_channel):
+        await await_n(channel.delete() for channel in self.get_pair_channels(ctx, meta_channel))
+        self.pair_channels = []
+
+    async def create_pair_channels(self, ctx, meta_channel, n):
+        await self.destroy_pair_channels(ctx, meta_channel)
+        futures = []
+        for i in range(1, n + 1):
+            futures.append(ctx.guild.create_voice_channel(str(i), category=meta_channel))
+        return await await_n(futures)
 
     @command(
         names = ('init', 'create', 'inti', 'craete', 'cretae', 'c', 'i', '+'),
@@ -56,7 +84,7 @@ class Client(CommandClient):
             or await self.create_category(ctx)
         )
         self.lobby_channel = (
-            self.get_lobby_channel(ctx)
+            self.get_lobby_channel(ctx, self.meta_channel)
             or await self.create_lobby(ctx)
         )
 
@@ -66,10 +94,83 @@ class Client(CommandClient):
     )
     async def destroy(self, ctx):
         futures = []
-        for channel in (self.get_meta_channel(ctx), self.get_lobby_channel(ctx)):
+        meta_channel = self.get_meta_channel(ctx)
+        for channel in (self.get_lobby_channel(ctx, meta_channel), meta_channel):
             if channel:
                 futures.append(channel.delete())
-        await asyncio.wait(futures)
+        await await_n(futures)
+        self.lobby_channel = None
+        self.meta_channel = None
+        self.pair_channels = []
+
+    async def get_channels(self, ctx):
+        meta_channel = self.get_meta_channel(ctx)
+        lobby_channel = self.get_lobby_channel(ctx, meta_channel)
+        if meta_channel is None or lobby_channel is None:
+            await ctx.answer('error: cannot start shuffling, you need to initialize channels')
+            await self.help(ctx)
+            return None
+        return meta_channel, lobby_channel
+
+    @command(
+        names = ('shuffle', 'start', 'run', 'strat', 'rnu'),
+        description = 'start shuffling'
+    )
+    async def shuffle(self, ctx):
+        channels = await self.get_channels(ctx)
+        if not channels: return
+        meta_channel, lobby_channel = channels
+        members = lobby_channel.members[:]
+        slots = len(members) >> 1
+        self.pair_channels = await self.create_pair_channels(ctx, meta_channel, slots)
+        slots = []
+        for i, _ in enumerate(self.pair_channels):
+            slots.append(i)
+            slots.append(i)
+        shuffle(slots)
+        futures = []
+        for slot in slots:
+            member = members.pop()
+            if member is None: break
+            futures.append(member.move_to(self.pair_channels[slot]))
+        if members:
+            futures.append(members.pop().move_to(choice(self.pair_channels)))
+        await await_n(futures)
+
+    @command(
+        names = ('stop', 'quit', 'exit', 'abort', 'back', 'return'),
+        description = 'move everyone back to lobby'
+    )
+    async def stop(self, ctx):
+        channels = await self.get_channels(ctx)
+        if not channels: return
+        meta_channel, lobby_channel = channels
+        pair_channels = self.get_pair_channels(ctx, meta_channel)
+        futures = []
+        for channel in pair_channels:
+            for member in channel.members:
+                futures.append(member.move_to(lobby_channel))
+        await await_n(futures)
+        await self.destroy_pair_channels(ctx, meta_channel)
+
+    @command(
+        names = ('loop',),
+        description = 'repeat "shuffle" and "stop" <n> (default: 3) times and <t> (default: 120) seconds'
+    )
+    async def loop(self, ctx):
+        if len(ctx.args) >= 1 and ctx.args[0].isdigit():
+            n = int(ctx.args[0])
+        else:
+            n = 3
+        if len(ctx.args) >= 2 and ctx.args[1].isdigit():
+            t = int(ctx.args[1])
+        else:
+            t = 120
+        await ctx.answer(f'repeat shuffling {n} times and each {t} seconds')
+        for _ in range(n):
+            await self.shuffle(ctx)
+            sleep(t)
+            await self.stop(ctx)
 
 
 if __name__ == '__main__':
@@ -78,5 +179,5 @@ if __name__ == '__main__':
     if token is None:
         print('error: no token was given')
         exit(1)
-    bot = Client(activity=discord.Game(name=config))
+    bot = Client(activity=discord.Game(name=config.GAME_STATUS))
     bot.run(token)
diff --git a/config.py b/config.py
index 404c66a..8ba8767 100644
--- a/config.py
+++ b/config.py
@@ -3,6 +3,8 @@ TOKEN_ENV_VAR = 'DISCORD_TOKEN'
 
 COMMAND_PREFIX = '!<3'
 
+GAME_STATUS = 'with love'
+
 HELP_TEXT = f'''{NAME.title()} - your partner for getting shuffled.
 {NAME.title()} is a discord bot to get to know your people.
 This software is open-source <https://git.kobert.dev/lovefinderrz.git/>!