diff --git a/Cargo.lock b/Cargo.lock index 3089c31..6910426 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,7 +4,7 @@ version = 4 [[package]] name = "RhythmWorks" -version = "0.5.3" +version = "0.6.0" dependencies = [ "color-eyre", "crossterm", @@ -24,6 +24,7 @@ dependencies = [ "songbird", "symphonia", "tokio", + "tokio-util", "tonic", "tracing", "tracing-opentelemetry", @@ -4980,9 +4981,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.17" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2efa149fe76073d6e8fd97ef4f4eca7b67f599660115591483572e406e165594" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" dependencies = [ "bytes", "futures-core", diff --git a/Cargo.toml b/Cargo.toml index eefe6ba..0cbe0db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "RhythmWorks" -version = "0.5.3" +version = "0.6.0" edition = "2024" [dependencies] @@ -22,6 +22,7 @@ serenity = { version = "0.12.5", features = ["client", "gateway", "voice"] } songbird = { version = "0.5.0", features = ["builtin-queue", "driver", "serenity"] } symphonia = { version = "0.5.5", features = ["aac", "alac", "isomp4", "mp3"] } tokio = "1.48.0" +tokio-util = "0.7.18" tonic = { version = "0.12.3", features = ["tls-roots"] } tracing = "0.1.41" tracing-opentelemetry = "0.32.0" diff --git a/src/main.rs b/src/main.rs index eb4a8eb..73a6806 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,8 +21,14 @@ use songbird::input::Input; use std::env; use std::process::Command; use std::process::Stdio; -pub mod modules; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; +pub mod modules; +enum BotState { + Stopped, + Running, +} struct TrackTraceHandler { otel_ctx: opentelemetry::Context, guild_id: String, @@ -349,6 +355,30 @@ impl EventHandler for Handler { println!("{} is connected!", ready.user.name); } } +async fn run_bot(token: String, shutdown: CancellationToken) { + let intents = GatewayIntents::GUILD_MESSAGES + | GatewayIntents::GUILDS + | GatewayIntents::DIRECT_MESSAGES + | GatewayIntents::MESSAGE_CONTENT + | GatewayIntents::GUILD_VOICE_STATES; + + let mut client = Client::builder(&token, intents) + .event_handler(Handler) + .register_songbird() + .await + .expect("Err creating client"); + + tokio::select! { + result = client.start() => { + if let Err(why) = result { + println!("Client error: {why:?}"); + } + } + _ = shutdown.cancelled() => { + println!("Shutting down bot..."); + } + } +} #[tokio::main] async fn main() { @@ -365,9 +395,18 @@ This is free software, and you are welcome to redistribute it under certain cond .await; let service_name = "RhythmWorks"; let username = env::var("USER").or_else(|_| env::var("USERNAME")).unwrap(); + let mut bot_task: Option> = None; + let mut shutdown_token: Option = None; + let mut bot_state = BotState::Stopped; loop { let token = credentialmanager::get_token(service_name, &username); - let options: Vec<&'static str> = vec!["- Start Bot", "- Change Token", "- Delete Token"]; + let options: Vec<&'static str> = vec![ + "- Start Bot", + "- Stop Bot", + "- Change Token", + "- Delete Token", + "- Exit", + ]; let selection = Select::with_theme(&ColorfulTheme::default()) .with_prompt("Choose an Option") .default(0) @@ -376,26 +415,44 @@ This is free software, and you are welcome to redistribute it under certain cond .unwrap(); match selection { 0 => { - let intents = GatewayIntents::GUILD_MESSAGES - | GatewayIntents::GUILDS - | GatewayIntents::DIRECT_MESSAGES - | GatewayIntents::MESSAGE_CONTENT - | GatewayIntents::GUILD_VOICE_STATES; - let mut client = Client::builder(&token, intents) - .event_handler(Handler) - .register_songbird() - .await - .expect("Err creating client"); - if let Err(why) = client.start().await { - println!("Client error: {why:?}"); + if matches!(bot_state, BotState::Running) { + println!("[x] Bot is already running."); + continue; } + + let token = credentialmanager::get_token(service_name, &username); + let shutdown = CancellationToken::new(); + let shutdown_clone = shutdown.clone(); + + let task = tokio::spawn(run_bot(token, shutdown_clone)); + + bot_task = Some(task); + shutdown_token = Some(shutdown); + bot_state = BotState::Running; + + println!("✅ Bot started."); } 1 => { - credentialmanager::change_token(service_name, &username); + if let Some(token) = shutdown_token.take() { + token.cancel(); + } + if let Some(task) = bot_task.take() { + let _ = task.await; + } + println!("Bot stopped."); } 2 => { + credentialmanager::change_token(service_name, &username); + } + 3 => { credentialmanager::remove_token(service_name, &username); } + 4 => { + if let Some(token) = shutdown_token { + token.cancel(); + } + break; + } _ => println!("Invalid Option"), } }