From 726bee497a8798c222ecacca6370deebe2aa5f5a Mon Sep 17 00:00:00 2001
From: trivernis <trivernis@protonmail.com>
Date: Thu, 22 Apr 2021 10:30:20 +0200
Subject: [PATCH] Fix handling of forced vc disconnect

Signed-off-by: trivernis <trivernis@protonmail.com>
---
 src/handler.rs                  | 40 ++++++++++++++++++++++++++-------
 src/providers/music/lavalink.rs | 11 ++++++++-
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/src/handler.rs b/src/handler.rs
index 99fa0e7..76389e9 100644
--- a/src/handler.rs
+++ b/src/handler.rs
@@ -1,6 +1,6 @@
 use serenity::async_trait;
 use serenity::client::Context;
-use serenity::model::channel::Reaction;
+use serenity::model::channel::{GuildChannel, Reaction};
 use serenity::model::event::ResumedEvent;
 use serenity::model::gateway::{Activity, Ready};
 use serenity::model::guild::Member;
@@ -9,6 +9,7 @@ use serenity::model::voice::VoiceState;
 use serenity::prelude::*;
 
 use crate::commands::music::get_music_player_for_guild;
+use crate::utils::context_data::MusicPlayers;
 use crate::utils::delete_messages_from_database;
 use bot_serenityutils::menu::{
     handle_message_delete, handle_message_delete_bulk, handle_reaction_add, handle_reaction_remove,
@@ -109,12 +110,12 @@ impl EventHandler for Handler {
             return;
         };
 
-        if let Some(old_id) = old_state.and_then(|c| c.channel_id) {
-            member_count = get_own_channel_member_count(&ctx, &old_id).await;
+        if let Some(old_id) = old_state.clone().and_then(|c| c.channel_id) {
+            member_count = get_own_channel_member_count(&ctx, old_id).await;
         }
         if member_count.is_none() {
             if let Some(new_id) = new_state.channel_id {
-                member_count = get_own_channel_member_count(&ctx, &new_id).await;
+                member_count = get_own_channel_member_count(&ctx, new_id).await;
             }
         }
 
@@ -126,14 +127,28 @@ impl EventHandler for Handler {
                 player.set_leave_flag(count == 0);
             }
         }
+        if let (Some(state), None) = (old_state, new_state.channel_id) {
+            let current_user = ctx.cache.current_user().await;
+
+            if state.user_id == current_user.id {
+                let mut data = ctx.data.write().await;
+                let players = data.get_mut::<MusicPlayers>().unwrap();
+
+                if let Some(player) = players.remove(&guild_id.0) {
+                    let mut player = player.lock().await;
+                    let _ = player.delete_now_playing().await;
+                    let _ = player.stop().await;
+                }
+            }
+        }
     }
 }
 
 /// Returns the number of members in the channel if it's the bots voice channel
-async fn get_own_channel_member_count(ctx: &Context, channel_id: &ChannelId) -> Option<usize> {
-    let channel = ctx.http.get_channel(channel_id.0).await.ok()?;
-    let guild_channel = channel.guild()?;
-    let current_user = ctx.http.get_current_user().await.ok()?;
+async fn get_own_channel_member_count(ctx: &Context, channel_id: ChannelId) -> Option<usize> {
+    let guild_channel = get_guild_channel(ctx, channel_id).await?;
+
+    let current_user = ctx.cache.current_user().await;
 
     let members = guild_channel.members(&ctx).await.ok()?;
     let own_channel = members
@@ -148,3 +163,12 @@ async fn get_own_channel_member_count(ctx: &Context, channel_id: &ChannelId) ->
 
     Some(members.len())
 }
+
+/// Returns the guild channel for a guild ID
+async fn get_guild_channel(ctx: &Context, channel_id: ChannelId) -> Option<GuildChannel> {
+    if let Some(channel) = ctx.cache.channel(channel_id).await {
+        return channel.guild();
+    }
+    let channel = ctx.http.get_channel(channel_id.0).await.ok()?;
+    channel.guild()
+}
diff --git a/src/providers/music/lavalink.rs b/src/providers/music/lavalink.rs
index ea10954..f5e9813 100644
--- a/src/providers/music/lavalink.rs
+++ b/src/providers/music/lavalink.rs
@@ -1,6 +1,6 @@
 use crate::utils::context_data::MusicPlayers;
 use lavalink_rs::gateway::LavalinkEventHandler;
-use lavalink_rs::model::{TrackFinish, TrackStart};
+use lavalink_rs::model::{PlayerUpdate, Stats, TrackFinish, TrackStart};
 use lavalink_rs::LavalinkClient;
 use serenity::async_trait;
 use serenity::prelude::TypeMapKey;
@@ -17,6 +17,7 @@ impl LavalinkEventHandler for LavalinkHandler {
     async fn track_start(&self, _client: LavalinkClient, event: TrackStart) {
         log::info!("Track started!\nGuild: {}", event.guild_id);
     }
+
     async fn track_finish(&self, _: LavalinkClient, event: TrackFinish) {
         log::info!("Track finished!\nGuild: {}", event.guild_id);
         let player = {
@@ -35,6 +36,14 @@ impl LavalinkEventHandler for LavalinkHandler {
             }
         }
     }
+
+    async fn player_update(&self, _: LavalinkClient, event: PlayerUpdate) {
+        log::debug!("Received player update event: {:?}", event);
+    }
+
+    async fn stats(&self, _: LavalinkClient, event: Stats) {
+        log::debug!("Received stats event: {:?}", event);
+    }
 }
 
 pub struct Lavalink;