Functional bot: responds to commands, moves users

src/main.rs

@@ -1 +1,188 @@
-fn main() {}
+use std::{
+    collections::BTreeMap,
+    sync::{Arc, RwLock},
+};
+
+use eyre::{bail, Context as _, Error, OptionExt, Result};
+use futures::FutureExt;
+use poise::serenity_prelude as serenity;
+use rand::seq::SliceRandom;
+
+// TODO: persist configuration
+// TODO: buttons?
+
+#[derive(Copy, Clone)]
+struct GuildData {
+    town_square: serenity::ChannelId,
+    cottages: serenity::ChannelId,
+}
+
+#[derive(Default)]
+struct Data(Arc<RwLock<BTreeMap<serenity::GuildId, GuildData>>>);
+
+type Context<'a> = poise::Context<'a, Data, Error>;
+
+async fn run(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<_>>();
+    users.shuffle(&mut rand::thread_rng());
+
+    let to = channels_for(to)?;
+
+    futures::future::try_join_all(
+        to.into_iter()
+            .cycle()
+            .zip(&users)
+            .map(|(channel, user)| guild.move_member(ctx, user.user.id, channel)),
+    )
+    .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
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+    let GuildData {
+        town_square,
+        cottages,
+    } = ctx
+        .data()
+        .0
+        .read()
+        .expect("No poisoned locks")
+        .get(&guild)
+        .copied()
+        .ok_or_eyre("This bot hasn't been configured for this server")?;
+
+    run(ctx, town_square, cottages).await?;
+
+    Ok(())
+}
+
+#[poise::command(slash_command, prefix_command, ephemeral)]
+async fn dawn(ctx: Context<'_>) -> Result<()> {
+    let guild = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+    let GuildData {
+        town_square,
+        cottages,
+    } = ctx
+        .data()
+        .0
+        .read()
+        .expect("No poisoned locks")
+        .get(&guild)
+        .copied()
+        .ok_or_eyre("This bot hasn't been configured for this server")?;
+
+    run(ctx, cottages, town_square).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,
+) -> Result<()> {
+    let guild = ctx
+        .guild_id()
+        .ok_or_eyre("This bot only works in servers")?;
+
+    ctx.data().0.write().expect("No poisoned locks").insert(
+        guild,
+        GuildData {
+            town_square,
+            cottages,
+        },
+    );
+    ctx.reply("Configured successfully!").await?;
+    Ok(())
+}
+
+async fn register_commands(
+    ctx: impl AsRef<serenity::Http>,
+    options: &poise::FrameworkOptions<Data, Error>,
+    guild: serenity::GuildId,
+) -> Result<()> {
+    poise::builtins::register_in_guild(ctx, &options.commands, guild).await?;
+    Ok(())
+}
+
+async fn on_event(
+    ctx: &serenity::Context,
+    event: &serenity::FullEvent,
+    framework: poise::FrameworkContext<'_, Data, Error>,
+    _: &Data,
+) -> Result<()> {
+    if let serenity::FullEvent::GuildCreate { guild, is_new: _ } = event {
+        register_commands(ctx, framework.options, guild.id).await?;
+    };
+    Ok(())
+}
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> Result<()> {
+    color_eyre::install()?;
+    let token = std::fs::read_to_string(
+        std::env::var_os("DISCORD_TOKEN_FILE").ok_or_eyre("A Discord token is required")?,
+    )
+    .wrap_err("A Discord token is required")?;
+    let intents = serenity::GatewayIntents::non_privileged();
+
+    let framework = poise::Framework::builder()
+        .options(poise::FrameworkOptions {
+            commands: vec![dusk(), dawn(), configure()],
+            event_handler: |ctx, ev, ctxf, data| on_event(ctx, ev, ctxf, data).boxed(),
+            ..Default::default()
+        })
+        .setup(|ctx, ready, framework| {
+            Box::pin(async move {
+                for guild in &ready.guilds {
+                    register_commands(ctx, framework.options(), guild.id).await?;
+                }
+                Ok(Data::default())
+            })
+        })
+        .build();
+
+    serenity::ClientBuilder::new(token, intents)
+        .framework(framework)
+        .await?
+        .start()
+        .await?;
+
+    Ok(())
+}