diff --git a/src/main.rs b/src/main.rs index 8761bd3..5a64a5b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,11 @@ use std::path::Path; -use eyre::{bail, Context as _, Error, OptionExt, Result}; +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, @@ -77,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> { + 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::>(); + 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)), ) @@ -123,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::>(); + + 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 @@ -134,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(()) } @@ -150,7 +169,7 @@ async fn dawn(ctx: Context<'_>) -> Result<()> { .. } = ctx.data().get(guild).await?; - run(ctx, cottages, town_square).await?; + join(ctx, cottages, town_square).await?; Ok(()) } @@ -163,7 +182,6 @@ async fn st(ctx: Context<'_>, mut spectators: Vec) -> Result<( let GuildData { st, .. } = ctx.data().get(guild).await?; spectators.push(ctx.author().id); - guild .members_iter(&ctx) .try_for_each(|member| { @@ -186,12 +204,21 @@ async fn st(ctx: Context<'_>, mut spectators: Vec) -> Result<( #[poise::command(slash_command, ephemeral)] async fn currently_playing(ctx: Context<'_>, players: Vec) -> Result<()> { - let guild = ctx.guild_id().ok_or_eyre("This bot only works in servers")?; - let GuildData { currently_playing, .. } = ctx.data().get(guild).await?; + 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?; + guild + .member(&ctx, player) + .await? + .add_role(&ctx, currently_playing) + .await?; } - ctx.reply(format!("Given {} members currently playing", players.len())).await?; + ctx.reply(format!("Given {} members currently playing", players.len())) + .await?; Ok(()) } @@ -201,7 +228,7 @@ async fn configure( #[channel_types("Voice")] town_square: serenity::ChannelId, #[channel_types("Category")] cottages: serenity::ChannelId, st: serenity::RoleId, - currently_playing: serenity::RoleId + currently_playing: serenity::RoleId, ) -> Result<()> { let guild = ctx .guild_id() @@ -250,7 +277,8 @@ 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() | serenity::GatewayIntents::GUILD_MEMBERS; + let intents = + serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::GUILD_MEMBERS; let framework = poise::Framework::builder() .options(poise::FrameworkOptions {