Core: asyncify choose and implement role filter

youmubot-core/src/community/mod.rs

@@ -9,6 +9,7 @@ use serenity::{
     },
     model::{
         channel::{Channel, Message},
+        id::RoleId,
         user::OnlineStatus,
     },
     utils::MessageBuilder,
@@ -30,39 +31,57 @@ struct Community;
 #[command]
 #[description = r"👑 Randomly choose an active member and mention them!
 Note that only online/idle users in the channel are chosen from."]
-#[usage = "[title = the chosen one]"]
+#[usage = "[title = the chosen one] / [limited roles = everyone online]"]
 #[example = "the strongest in Gensokyo"]
 #[bucket = "community"]
-#[max_args(1)]
+#[max_args(2)]
-pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
+pub async fn choose(ctx: &Context, m: &Message, mut args: Args) -> CommandResult {
     let title = if args.is_empty() {
         "the chosen one".to_owned()
     } else {
         args.single::<String>()?
     };
+    let role = match args.single::<RoleId>().ok() {
+        Some(v) => v.to_role_cached(&ctx).await,
+        None => None,
+    };
 
     let users: Result<Vec<_>, Error> = {
-        let guild = m.guild(&ctx).unwrap();
+        let guild = m.guild(&ctx).await.unwrap();
-        let guild = guild.read();
         let presences = &guild.presences;
-        let channel = m.channel_id.to_channel(&ctx)?;
+        let channel = m.channel_id.to_channel(&ctx).await?;
         if let Channel::Guild(channel) = channel {
-            let channel = channel.read();
             Ok(channel
-                .members(&ctx)?
+                .members(&ctx)
+                .await?
                 .into_iter()
-                .filter(|v| !v.user.read().bot)
+                .filter(|v| !v.user.bot) // Filter out bots
-                .map(|v| v.user_id())
                 .filter(|v| {
+                    // Filter out only online people
                     presences
-                        .get(v)
+                        .get(&v.user.id)
                         .map(|presence| {
                             presence.status == OnlineStatus::Online
                                 || presence.status == OnlineStatus::Idle
                         })
                         .unwrap_or(false)
                 })
-                .collect())
+                .map(|mem| future::ready(mem))
+                .collect::<stream::FuturesUnordered<_>>()
+                .filter(|member| async {
+                    // Filter by role if provided
+                    if let Some(role) = role {
+                        member
+                            .roles(&ctx)
+                            .await
+                            .map(|roles| roles.into_iter().any(|r| role.id == r.id))
+                            .unwrap_or(false)
+                    } else {
+                        true
+                    }
+                })
+                .collect()
+                .await)
         } else {
             panic!()
         }
@@ -73,7 +92,8 @@ pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
         m.reply(
             &ctx,
             "🍰 Have this cake for yourself because no-one is here for the gods to pick.",
-        )?;
+        )
+        .await?;
         return Ok(());
     }
 
@@ -83,19 +103,26 @@ pub fn choose(ctx: &mut Context, m: &Message, mut args: Args) -> CommandResult {
         &users[uniform.sample(&mut rng)]
     };
 
-    m.channel_id.send_message(&ctx, |c| {
+    m.channel_id
+        .send_message(&ctx, |c| {
             c.content(
                 MessageBuilder::new()
                     .push("👑 The Gensokyo gods have gathered around and decided, out of ")
                     .push_bold(format!("{}", users.len()))
-                .push(" potential prayers, ")
+                    .push(" ")
+                    .push(
+                        role.map(|r| r.mention() + "s")
+                            .unwrap_or("potential prayers".to_owned()),
+                    )
+                    .push(", ")
                     .push(winner.mention())
                     .push(" will be ")
                     .push_bold_safe(title)
                     .push(". Congrats! 🎉 🎊 🥳")
                     .build(),
             )
-    })?;
+        })
+        .await?;
 
     Ok(())
 }