Refactor to allow only spreading into empty channels

Split `run` into `spread` and `join`

src/main.rs

@@ -1,18 +1,18 @@
 use std::path::Path;
 
-use eyre::{bail, Context as _, Error, OptionExt, Result};
-use futures::FutureExt;
+use eyre::{Context as _, Error, OptionExt, Result};
+use futures::{FutureExt, TryStreamExt};
 use poise::serenity_prelude as serenity;
 use rand::seq::SliceRandom;
 use rusqlite::OptionalExtension;
 
-// TODO: buttons?
-
 #[derive(Copy, Clone)]
 struct GuildData {
     guild: serenity::GuildId,
     town_square: serenity::ChannelId,
     cottages: serenity::ChannelId,
+    st: serenity::RoleId,
+    currently_playing: serenity::RoleId,
 }
 
 struct Data(tokio::sync::Mutex<rusqlite::Connection>);
@@ -20,7 +20,7 @@ struct Data(tokio::sync::Mutex<rusqlite::Connection>);
 impl Data {
     fn open(path: impl AsRef<Path>) -> Result<Self> {
         let conn = rusqlite::Connection::open(path)?;
-        conn.execute_batch("CREATE TABLE IF NOT EXISTS guilds(guild TEXT PRIMARY KEY NOT NULL, town_square TEXT NOT NULL, cottages TEXT NOT NULL);")?;
+        conn.execute_batch("CREATE TABLE IF NOT EXISTS guilds(guild TEXT PRIMARY KEY NOT NULL, town_square TEXT NOT NULL, cottages TEXT NOT NULL, st TEXT, currently_playing TEXT);")?;
         Ok(Self(tokio::sync::Mutex::new(conn)))
     }
 
@@ -28,12 +28,14 @@ impl Data {
         self.0
             .lock()
             .await
-            .prepare_cached("SELECT guild, town_square, cottages FROM guilds WHERE guild = ?;")?
+            .prepare_cached("SELECT guild, town_square, cottages, st, currently_playing FROM guilds WHERE guild = ?;")?
             .query_row(rusqlite::params![guild.to_string()], |row| {
                 Ok(GuildData {
                     guild: row.get_unwrap::<_, String>("guild").parse().unwrap(),
                     town_square: row.get_unwrap::<_, String>("town_square").parse().unwrap(),
                     cottages: row.get_unwrap::<_, String>("cottages").parse().unwrap(),
+                    st: row.get_unwrap::<_, String>("st").parse().unwrap(),
+                    currently_playing: row.get_unwrap::<_, String>("currently_playing").parse().unwrap(),
                 })
             })
             .optional()?
@@ -46,16 +48,26 @@ impl Data {
             guild,
             town_square,
             cottages,
+            st,
+            currently_playing,
         }: GuildData,
     ) -> Result<()> {
         self.0
             .lock()
             .await
-            .prepare_cached("INSERT INTO guilds (guild, town_square, cottages) VALUES (?, ?, ?) ON CONFLICT(guild) DO UPDATE SET town_square = excluded.town_square, cottages = excluded.cottages")?
+            .prepare_cached("INSERT INTO guilds (guild, town_square, cottages, st, currently_playing) VALUES (?, ?, ?, ?, ?)
+                 ON CONFLICT(guild) DO UPDATE SET
+                     town_square = excluded.town_square,
+                     cottages = excluded.cottages,
+                     st = excluded.st,
+                     currently_playing = excluded.currently_playing;
+            ")?
             .execute(rusqlite::params![
                 guild.to_string(),
                 town_square.to_string(),
-                cottages.to_string()
+                cottages.to_string(),
+                st.to_string(),
+                currently_playing.to_string(),
             ])?;
         Ok(())
     }
@@ -63,43 +75,45 @@ impl Data {
 
 type Context<'a> = poise::Context<'a, Data, Error>;
 
-async fn run(ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId) -> Result<()> {
+async fn channel_children(
+    ctx: &Context<'_>,
+    channel: serenity::ChannelId,
+) -> Result<Vec<serenity::GuildChannel>> {
+    let channels = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?
+        .channels(&ctx)
+        .await?;
+    Ok(channels
+        .into_values()
+        .filter(|chan| chan.parent_id == Some(channel))
+        .collect())
+}
+
+async fn spread(
+    ctx: Context<'_>,
+    from: serenity::ChannelId,
+    to: serenity::ChannelId,
+) -> Result<()> {
     let guild = ctx
         .guild_id()
         .ok_or_eyre("This bot only works in servers")?;
 
-    let channels = guild.channels(&ctx).await?;
-    let channels_for = |id| {
-        let channel = channels
-            .get(&id)
-            .ok_or_eyre(format!("The channel {id:?} doesn't exist"))?;
-        Ok::<_, Error>(match channel.kind {
-            serenity::ChannelType::Voice => vec![channel],
-            serenity::ChannelType::Category => channels
-                .values()
-                .filter(|chan| {
-                    chan.kind == serenity::ChannelType::Voice && chan.parent_id == Some(channel.id)
-                })
-                .collect(),
-            ty => bail!(
-                "This bot doesn't work with non-voice channels: {:?} has type {:?}",
-                &channel.name,
-                ty
-            ),
-        })
-    };
-    let mut users = channels_for(from)?
-        .into_iter()
-        .flat_map(|chan| chan.members(ctx).unwrap())
-        .collect::<Vec<_>>();
+    let from = from
+        .to_channel(ctx)
+        .await?
+        .guild()
+        .ok_or_eyre("This bot only works in servers")?;
+    let mut to = channel_children(&ctx, to).await?;
+
+    let mut users = from.members(ctx)?;
     users.shuffle(&mut rand::thread_rng());
 
-    let mut to = channels_for(to)?;
+    to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty()));
     to.shuffle(&mut rand::thread_rng());
 
     futures::future::try_join_all(
         to.into_iter()
-            .cycle()
             .zip(&users)
             .map(|(channel, user)| guild.move_member(ctx, user.user.id, channel)),
     )
@@ -109,6 +123,25 @@ async fn run(ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelI
     Ok(())
 }
 
+async fn join(ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId) -> Result<()> {
+    let guild = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+
+    let users = channel_children(&ctx, from)
+        .await?
+        .into_iter()
+        .flat_map(|x| x.members(ctx))
+        .flatten()
+        .collect::<Vec<_>>();
+
+    futures::future::try_join_all(users.iter().map(|user| guild.move_member(&ctx, user, to)))
+        .await?;
+
+    ctx.reply(format!("Moved {} users", users.len())).await?;
+    Ok(())
+}
+
 #[poise::command(slash_command, prefix_command, ephemeral)]
 async fn dusk(ctx: Context<'_>) -> Result<()> {
     let guild = ctx
@@ -120,7 +153,7 @@ async fn dusk(ctx: Context<'_>) -> Result<()> {
         ..
     } = ctx.data().get(guild).await?;
 
-    run(ctx, town_square, cottages).await?;
+    spread(ctx, town_square, cottages).await?;
 
     Ok(())
 }
@@ -136,26 +169,81 @@ async fn dawn(ctx: Context<'_>) -> Result<()> {
         ..
     } = ctx.data().get(guild).await?;
 
-    run(ctx, cottages, town_square).await?;
+    join(ctx, cottages, town_square).await?;
 
     Ok(())
 }
 
+#[poise::command(slash_command, ephemeral)]
+async fn st(ctx: Context<'_>, mut spectators: Vec<serenity::UserId>) -> Result<()> {
+    let guild = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+    let GuildData { st, .. } = ctx.data().get(guild).await?;
+    spectators.push(ctx.author().id);
+
+    guild
+        .members_iter(&ctx)
+        .try_for_each(|member| {
+            let spectators = &spectators;
+            async move {
+                if spectators.contains(&member.user.id) {
+                    member.add_role(&ctx, st).await
+                } else {
+                    member.remove_role(&ctx, st).await
+                }
+            }
+        })
+        .await?;
+
+    ctx.reply(format!("{} members given ST", spectators.len()))
+        .await?;
+
+    Ok(())
+}
+
+#[poise::command(slash_command, ephemeral)]
+async fn currently_playing(ctx: Context<'_>, players: Vec<serenity::UserId>) -> Result<()> {
+    let guild = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+    let GuildData {
+        currently_playing, ..
+    } = ctx.data().get(guild).await?;
+    for player in &players {
+        guild
+            .member(&ctx, player)
+            .await?
+            .add_role(&ctx, currently_playing)
+            .await?;
+    }
+    ctx.reply(format!("Given {} members currently playing", players.len()))
+        .await?;
+    Ok(())
+}
+
 #[poise::command(slash_command, prefix_command, ephemeral)]
 async fn configure(
     ctx: Context<'_>,
     #[channel_types("Voice")] town_square: serenity::ChannelId,
     #[channel_types("Category")] cottages: serenity::ChannelId,
+    st: serenity::RoleId,
+    currently_playing: serenity::RoleId,
 ) -> Result<()> {
     let guild = ctx
         .guild_id()
         .ok_or_eyre("This bot only works in servers")?;
 
-    ctx.data().insert(GuildData {
+    ctx.data()
+        .insert(GuildData {
             guild,
             town_square,
             cottages,
-    }).await?;
+            st,
+            currently_playing,
+        })
+        .await?;
+
     ctx.reply("Configured successfully!").await?;
     Ok(())
 }
@@ -189,11 +277,12 @@ async fn main() -> Result<()> {
     )
     .wrap_err("A Discord token is required")?;
     let dbfile = std::env::var_os("DB_FILE").ok_or_eyre("A database file is required")?;
-    let intents = serenity::GatewayIntents::non_privileged();
+    let intents =
+        serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::GUILD_MEMBERS;
 
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {
-            commands: vec![dusk(), dawn(), configure()],
+            commands: vec![dusk(), dawn(), st(), currently_playing(), configure()],
             event_handler: |ctx, ev, ctxf, data| on_event(ctx, ev, ctxf, data).boxed(),
             ..Default::default()
         })