Add equalizer command

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -209,7 +209,7 @@ dependencies = [
 
 [[package]]
 name = "bot-serenityutils"
-version = "0.2.3"
+version = "0.2.4"
 dependencies = [
  "futures",
  "log 0.4.14",
@@ -1095,6 +1095,7 @@ dependencies = [
  "serde",
  "serde-aux",
  "serde_json",
+ "serenity",
  "songbird",
  "tokio",
  "tokio-native-tls",

Cargo.toml

@@ -39,4 +39,4 @@ rustc_version_runtime = "0.2.0"
 trigram = "0.4.4"
 typemap_rev = "0.1.5"
 youtube-metadata = "0.1.1"
-lavalink-rs = {version="0.7.1", features=["native"]}
+lavalink-rs = {version="0.7.1", features=["native", "serenity"]}

bot-serenityutils/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "bot-serenityutils"
-version = "0.2.3"
+version = "0.2.4"
 authors = ["trivernis <trivernis@protonmail.com>"]
 edition = "2018"
 

bot-serenityutils/src/menu/controls.rs

@@ -124,7 +124,7 @@ pub async fn toggle_help(
 }
 
 /// Displays the menu page
-async fn display_page(ctx: &Context, menu: &mut Menu<'_>) -> SerenityUtilsResult<()> {
+pub async fn display_page(ctx: &Context, menu: &mut Menu<'_>) -> SerenityUtilsResult<()> {
     log::debug!("Displaying page {}", menu.current_page);
     let page = menu
         .pages

src/commands/music/equalizer.rs

@@ -0,0 +1,33 @@
+use serenity::framework::standard::macros::command;
+use serenity::framework::standard::{CommandError, CommandResult};
+use serenity::model::channel::Message;
+use serenity::prelude::*;
+
+use crate::commands::common::handle_autodelete;
+use crate::commands::music::{get_music_player_for_guild, DJ_CHECK};
+use crate::messages::music::equalizer::create_equalizer_message;
+use crate::messages::music::no_voicechannel::create_no_voicechannel_message;
+
+#[command]
+#[only_in(guilds)]
+#[description("Displays the equalizer for the music player")]
+#[usage("")]
+#[bucket("general")]
+#[checks(DJ)]
+async fn equalizer(ctx: &Context, msg: &Message) -> CommandResult {
+    let guild = msg.guild(&ctx.cache).await.unwrap();
+    log::debug!("Pausing playback for guild {}", guild.id);
+
+    let player = if let Some(player) = get_music_player_for_guild(ctx, guild.id).await {
+        player
+    } else {
+        return create_no_voicechannel_message(&ctx.http, msg.channel_id)
+            .await
+            .map_err(CommandError::from);
+    };
+
+    create_equalizer_message(&ctx, msg.channel_id, player).await?;
+    handle_autodelete(ctx, msg).await?;
+
+    Ok(())
+}

src/commands/music/mod.rs

@@ -11,22 +11,6 @@ use serenity::model::user::User;
 use songbird::Songbird;
 use tokio::sync::Mutex;
 
-use clear_queue::CLEAR_QUEUE_COMMAND;
-use current::CURRENT_COMMAND;
-use join::JOIN_COMMAND;
-use leave::LEAVE_COMMAND;
-use lyrics::LYRICS_COMMAND;
-use move_song::MOVE_SONG_COMMAND;
-use pause::PAUSE_COMMAND;
-use play::PLAY_COMMAND;
-use play_next::PLAY_NEXT_COMMAND;
-use playlists::PLAYLISTS_COMMAND;
-use queue::QUEUE_COMMAND;
-use remove_song::REMOVE_SONG_COMMAND;
-use save_playlist::SAVE_PLAYLIST_COMMAND;
-use shuffle::SHUFFLE_COMMAND;
-use skip::SKIP_COMMAND;
-
 use crate::providers::music::player::MusicPlayer;
 use crate::providers::music::queue::Song;
 use crate::providers::music::{add_youtube_song_to_database, youtube_dl};
@@ -41,6 +25,7 @@ use youtube_metadata::get_video_information;
 
 mod clear_queue;
 mod current;
+mod equalizer;
 mod join;
 mod leave;
 mod lyrics;
@@ -55,6 +40,23 @@ mod save_playlist;
 mod shuffle;
 mod skip;
 
+use clear_queue::CLEAR_QUEUE_COMMAND;
+use current::CURRENT_COMMAND;
+use equalizer::EQUALIZER_COMMAND;
+use join::JOIN_COMMAND;
+use leave::LEAVE_COMMAND;
+use lyrics::LYRICS_COMMAND;
+use move_song::MOVE_SONG_COMMAND;
+use pause::PAUSE_COMMAND;
+use play::PLAY_COMMAND;
+use play_next::PLAY_NEXT_COMMAND;
+use playlists::PLAYLISTS_COMMAND;
+use queue::QUEUE_COMMAND;
+use remove_song::REMOVE_SONG_COMMAND;
+use save_playlist::SAVE_PLAYLIST_COMMAND;
+use shuffle::SHUFFLE_COMMAND;
+use skip::SKIP_COMMAND;
+
 #[group]
 #[commands(
     join,
@@ -71,7 +73,8 @@ mod skip;
     playlists,
     lyrics,
     move_song,
-    remove_song
+    remove_song,
+    equalizer
 )]
 pub struct Music;
 

src/messages/music/equalizer.rs

@@ -0,0 +1,202 @@
+use crate::providers::music::player::MusicPlayer;
+use crate::utils::error::BotResult;
+use bot_serenityutils::core::EXTRA_LONG_TIMEOUT;
+use bot_serenityutils::error::SerenityUtilsResult;
+use bot_serenityutils::menu::{display_page, Menu, MenuBuilder, Page};
+use serenity::builder::{CreateEmbed, CreateMessage};
+use serenity::client::Context;
+use serenity::model::channel::Reaction;
+use serenity::model::id::ChannelId;
+use std::sync::atomic::{AtomicU8, Ordering};
+use std::sync::Arc;
+use tokio::sync::Mutex;
+use typemap_rev::TypeMapKey;
+
+static NEXT_BAND_BUTTON: &str = "➡️";
+static PREVIOUS_BAND_BUTTON: &str = "⬅️";
+static ADD_BUTTON: &str = "➕";
+static SUB_BUTTON: &str = "➖";
+
+struct SelectedBand;
+
+impl TypeMapKey for SelectedBand {
+    type Value = Arc<AtomicU8>;
+}
+
+struct Player;
+
+impl TypeMapKey for Player {
+    type Value = Arc<Mutex<MusicPlayer>>;
+}
+
+/// Creates a new equalizer message
+pub async fn create_equalizer_message(
+    ctx: &Context,
+    channel_id: ChannelId,
+    player: Arc<Mutex<MusicPlayer>>,
+) -> BotResult<()> {
+    let selected_band = Arc::new(AtomicU8::new(0));
+    let selected_band_clone = Arc::clone(&selected_band);
+    let player_clone = Arc::clone(&player);
+
+    MenuBuilder::default()
+        .add_page(Page::new_builder(move || {
+            let player = Arc::clone(&player_clone);
+            let selected_band = Arc::clone(&selected_band_clone);
+            Box::pin(async move {
+                let mut page = CreateMessage::default();
+                let mut embed = CreateEmbed::default();
+                create_equalizer_embed(selected_band.load(Ordering::Relaxed), &mut embed, &player)
+                    .await;
+
+                page.embed(|e| {
+                    e.0.clone_from(&embed.0);
+                    e
+                });
+
+                Ok(page)
+            })
+        }))
+        .add_control(0, PREVIOUS_BAND_BUTTON, |c, m, r| {
+            Box::pin(previous_band(c, m, r))
+        })
+        .add_help(PREVIOUS_BAND_BUTTON, "Selects the previous band.")
+        .add_control(1, NEXT_BAND_BUTTON, |c, m, r| Box::pin(next_band(c, m, r)))
+        .add_help(NEXT_BAND_BUTTON, "Selects the next band.")
+        .add_control(3, ADD_BUTTON, |c, m, r| Box::pin(add_to_band(c, m, r)))
+        .add_help(ADD_BUTTON, "Adds to the selected band.")
+        .add_control(2, SUB_BUTTON, |c, m, r| {
+            Box::pin(subtract_from_band(c, m, r))
+        })
+        .add_help(SUB_BUTTON, "Subtracts from the selected band")
+        .show_help()
+        .add_data::<SelectedBand>(selected_band)
+        .add_data::<Player>(player)
+        .timeout(EXTRA_LONG_TIMEOUT)
+        .build(ctx, channel_id)
+        .await?;
+    Ok(())
+}
+
+/// Creates a new equalizer embed
+async fn create_equalizer_embed<'a>(
+    selected_band: u8,
+    embed: &'a mut CreateEmbed,
+    player: &Arc<Mutex<MusicPlayer>>,
+) -> &'a mut CreateEmbed {
+    let mut description = String::new();
+    let bands = {
+        let player = player.lock().await;
+        player.get_equalizer().clone()
+    };
+    for i in 0..bands.len() {
+        if i as u8 == selected_band {
+            description += "⤋"
+        } else {
+            description += " ";
+        }
+    }
+    description += "\n";
+    for i in (0..11).rev() {
+        let eq_value = (i as f64) / 20.0 - 0.25;
+
+        for band in &bands {
+            if (eq_value > 0. && band >= &eq_value) || (eq_value < 0. && band <= &eq_value) {
+                description += "█";
+            } else if eq_value == 0. {
+                description += format!("-").as_str();
+            } else {
+                description += " ";
+            }
+        }
+        description += "\n";
+    }
+    for i in 0..bands.len() {
+        if i as u8 == selected_band {
+            description += "⤊"
+        } else {
+            description += " ";
+        }
+    }
+    embed
+        .title("Equalizer")
+        .description(format!("```\n{}\n```", description));
+
+    embed
+}
+
+/// Selects the previous band
+async fn next_band(ctx: &Context, menu: &mut Menu<'_>, _: Reaction) -> SerenityUtilsResult<()> {
+    let selected_band = menu.data.get::<SelectedBand>().unwrap();
+    if selected_band.load(Ordering::SeqCst) >= 14 {
+        selected_band.store(0, Ordering::SeqCst);
+    } else {
+        selected_band.fetch_add(1, Ordering::SeqCst);
+    }
+    display_page(ctx, menu).await?;
+
+    Ok(())
+}
+
+/// Selects the previous band
+async fn previous_band(ctx: &Context, menu: &mut Menu<'_>, _: Reaction) -> SerenityUtilsResult<()> {
+    let selected_band = menu.data.get::<SelectedBand>().unwrap();
+    if selected_band.load(Ordering::SeqCst) <= 0 {
+        selected_band.store(14, Ordering::SeqCst);
+    } else {
+        selected_band.fetch_sub(1, Ordering::SeqCst);
+    }
+    display_page(ctx, menu).await?;
+
+    Ok(())
+}
+
+/// Adds to the selected band
+async fn add_to_band(ctx: &Context, menu: &mut Menu<'_>, _: Reaction) -> SerenityUtilsResult<()> {
+    {
+        let selected_band = menu
+            .data
+            .get::<SelectedBand>()
+            .unwrap()
+            .load(Ordering::Relaxed);
+        let player = menu.data.get::<Player>().unwrap();
+        let mut player = player.lock().await;
+        let equalizer = player.get_equalizer();
+        let current_value = equalizer[selected_band as usize];
+
+        if current_value < 0.25 {
+            player.equalize(selected_band, current_value + 0.05).await?;
+        }
+    }
+
+    display_page(ctx, menu).await?;
+
+    Ok(())
+}
+
+/// Substracts from the selected band
+async fn subtract_from_band(
+    ctx: &Context,
+    menu: &mut Menu<'_>,
+    _: Reaction,
+) -> SerenityUtilsResult<()> {
+    {
+        let selected_band = menu
+            .data
+            .get::<SelectedBand>()
+            .unwrap()
+            .load(Ordering::Relaxed);
+        let player = menu.data.get::<Player>().unwrap();
+        let mut player = player.lock().await;
+        let equalizer = player.get_equalizer();
+        let current_value = equalizer[selected_band as usize];
+
+        if current_value > -0.25 {
+            player.equalize(selected_band, current_value - 0.05).await?;
+        }
+    }
+
+    display_page(ctx, menu).await?;
+
+    Ok(())
+}

src/messages/music/mod.rs

@@ -1,3 +1,4 @@
+pub mod equalizer;
 pub mod no_voicechannel;
 pub mod now_playing;
 pub mod queue;

src/providers/music/player.rs

@@ -3,7 +3,7 @@ use crate::providers::music::lavalink::Lavalink;
 use crate::providers::music::lyrics::get_lyrics;
 use crate::providers::music::queue::MusicQueue;
 use crate::utils::context_data::MusicPlayers;
-use crate::utils::error::BotResult;
+use crate::utils::error::{BotError, BotResult};
 use bot_serenityutils::core::{MessageHandle, SHORT_TIMEOUT};
 use bot_serenityutils::ephemeral_message::EphemeralMessage;
 use lavalink_rs::LavalinkClient;
@@ -28,6 +28,7 @@ pub struct MusicPlayer {
     msg_channel: ChannelId,
     leave_flag: bool,
     paused: bool,
+    equalizer: [f64; 15],
 }
 
 impl MusicPlayer {
@@ -47,6 +48,7 @@ impl MusicPlayer {
             now_playing_msg: None,
             leave_flag: false,
             paused: false,
+            equalizer: [0f64; 15],
         }
     }
 
@@ -232,6 +234,27 @@ impl MusicPlayer {
 
         Ok(())
     }
+
+    /// Returns the equalizer
+    pub fn get_equalizer(&self) -> &[f64; 15] {
+        &self.equalizer
+    }
+
+    /// Equalizes a specified band
+    pub async fn equalize(&mut self, band: u8, value: f64) -> BotResult<()> {
+        if band > 15 {
+            return Err(BotError::from("Invalid Equalizer band"));
+        }
+        if value < -0.25 || value > 0.25 {
+            return Err(BotError::from("Invalid Equalizer value"));
+        }
+        self.equalizer[band as usize] = value;
+        self.client
+            .equalize_all(self.guild_id, self.equalizer)
+            .await?;
+
+        Ok(())
+    }
 }
 
 /// Stats a tokio coroutine to check for player disconnect conditions