1 files changed, 104 insertions, 16 deletions

diff --git a/src/util/mod.rs b/src/util/mod.rs
index 9355e48..38f4e55 100644
--- a/src/util/mod.rs
+++ b/src/util/mod.rs
@@ -1,6 +1,19 @@
-use std::process::Stdio;
+use std::{
+    cmp::max_by,
+    collections::HashMap,
+    process::{
+        id,
+        Stdio,
+    },
+};
 
+use anyhow::anyhow;
 use chrono::Duration;
+use diesel::row::NamedRow;
+use fnv::{
+    FnvHashMap,
+    FnvHashSet,
+};
 use grate::tracing;
 use lazy_static::lazy_static;
 use poise::{
@@ -13,10 +26,14 @@ use regex::{
 };
 use serenity::{
     all::{
+        ChannelType,
         CreateMessage,
+        GuildId,
+        GuildRef,
         Message,
         Reaction,
         ReactionType,
+        VoiceState,
     },
     client::Context,
     model::{
@@ -27,6 +44,7 @@ use serenity::{
         permissions::Permissions,
     },
 };
+use tap::Pipe;
 use url::Url;
 
 use crate::{
@@ -49,29 +67,25 @@ pub async fn currently_playing(ctx: PoiseContext<'_>) -> bool {
     call.queue().current().is_some()
 }
 
-pub async fn users_listening(ctx: &Context) -> anyhow::Result<bool> {
-    let channel = CONFIG.discord.voice_channel().to_channel(&ctx).await?;
+pub async fn users_listening(ctx: PoiseContext<'_>) -> anyhow::Result<bool> {
+    let Some(channel) = best_voice_channel(ctx) else {
+        return Ok(false);
+    };
 
-    let res = channel
-        .guild()
-        .and_then(|ch| ch.guild(&ctx))
-        .map(|g| {
-            g.voice_states
-                .iter()
-                .any(|(_, state)| state.channel_id == Some(CONFIG.discord.voice_channel()))
-        })
-        .unwrap_or(false);
+    let Some(guild) = ctx.guild() else {
+        return Ok(false);
+    };
 
-    Ok(res)
+    guild.voice_states.iter().any(|(_, state)| state.channel_id == Some(channel)).pipe(Ok)
 }
 
 #[inline]
 pub fn msg<U, E>(ctx: poise::Context<'_, U, E>) -> Option<&Message> {
     match ctx {
         poise::Context::Prefix(poise::PrefixContext {
-            msg,
-            ..
-        }) => Some(msg),
+                                   msg,
+                                   ..
+                               }) => Some(msg),
         _ => None,
     }
 }
@@ -154,6 +168,80 @@ pub async fn unreact(
     Ok(())
 }
 
+#[inline]
+pub fn guild_id(ctx: PoiseContext<'_>) -> anyhow::Result<GuildId> {
+    ctx.guild_id().ok_or_else(|| anyhow::anyhow!("not in guild"))
+}
+
+#[inline]
+pub fn guild(ctx: PoiseContext<'_>) -> anyhow::Result<GuildRef<'_>> {
+    ctx.guild().ok_or_else(|| anyhow::anyhow!("not in guild"))
+}
+
+#[inline]
+pub fn author_voice_state(ctx: PoiseContext<'_>) -> Option<(VoiceState, GuildRef)> {
+    let guild = ctx.guild()?;
+    let caller_voice = guild.voice_states.get(&ctx.author().id)?.clone();
+
+    Some((caller_voice, guild))
+}
+
+#[inline]
+pub fn author_voice_channel(ctx: PoiseContext<'_>) -> Option<ChannelId> {
+    let (vs, _guild) = author_voice_state(ctx)?;
+    vs.channel_id
+}
+
+pub fn voice_states_by_channel(ctx: PoiseContext<'_>) -> HashMap<ChannelId, Vec<VoiceState>> {
+    let Some(guild) = ctx.guild() else {
+        return Default::default();
+    };
+
+    guild
+        .voice_states
+        .values()
+        .cloned()
+        .filter_map(|x| {
+            let id = x.channel_id?;
+            Some((id, x))
+        })
+        .fold(HashMap::new(), |mut acc, (id, state)| {
+            acc.entry(id).or_insert_with(|| vec![]).push(state);
+
+            acc
+        })
+}
+
+/// Select the most relevant voice channel for a given poise context.
+///
+/// - If the message's author is in a voice channel, use that.
+/// - If not, pick the most populated channel (channel age tiebreaks).
+pub fn best_voice_channel(ctx: PoiseContext<'_>) -> Option<ChannelId> {
+    if let Some(channel) = author_voice_channel(ctx) {
+        return Some(channel);
+    }
+
+    let voice_states = voice_states_by_channel(ctx);
+    let max_pop = voice_states.iter().map(|(_, states)| states.len()).max();
+
+    let matching_channels = voice_states
+        .iter()
+        .filter_map(|(&channel, state)| {
+            if state.len() == max_pop? {
+                return Some(channel);
+            }
+
+            None
+        })
+        .collect::<Vec<_>>();
+
+    if matching_channels.len() == 1 {
+        return matching_channels.first().cloned();
+    }
+
+    matching_channels.into_iter().min()
+}
+
 pub async fn send_result(
     ctx: &Context,
     channel: ChannelId,