Slight chat/play refactor, add roll command

service/src/bot/commands/audio/mod.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use log::{debug, warn};
 
+use reqwest::Url;
 use serenity::client::Cache;
 use serenity::model::application::interaction::{InteractionResponseType, application_command::ApplicationCommandInteraction};
 use serenity::model::prelude::{GuildId, ChannelId};
@@ -11,6 +12,8 @@ use siren::ServiceError;
 use songbird::{Call, Songbird};
 use songbird::input::{Restartable, Input, Metadata, error::Error as SongbirdError};
 
+use crate::bot::ytdlp::{PlaylistItem, YtDlp};
+
 pub mod pause;
 pub mod play;
 pub mod resume;
@@ -88,11 +91,7 @@ pub async fn edit_response(ctx: &Context, command: &ApplicationCommandInteractio
 }
 
 pub async fn add_song(call: Arc<Mutex<Call>>, url: &str, lazy: bool, volume: Option<f32>) -> Result<Metadata, SongbirdError> {
-  let source = if is_valid_url(url) {
-    Restartable::ytdl(url.to_owned(), lazy).await?
-  } else {
-    Restartable::ytdl_search(url, lazy).await?
-  };
+  let source = Restartable::ytdl(url.to_owned(), lazy).await?;
   let mut handler = call.lock().await;
   let track: Input = source.into();
   let metadata = *track.metadata.clone();
@@ -103,18 +102,41 @@ pub async fn add_song(call: Arc<Mutex<Call>>, url: &str, lazy: bool, volume: Opt
   Ok(metadata)
 }
 
-pub fn get_playlist_urls(url: &str) -> Result<Vec<String>, String> {
-  let mut urls: Vec<String> = Vec::new();
-  // TODO fix this later
-  urls.push(url.to_string());
-  Ok(urls)
+pub fn get_playlist_urls(url: &str) -> Result<Vec<PlaylistItem>, ServiceError> {
+  let output = YtDlp::new()
+    .arg("--flat-playlist")
+    .arg("--dump-json")
+    .arg(url)
+    .execute()?;
+  let items: Vec<PlaylistItem> = String::from_utf8(output.stdout)?
+    .split('\n')
+    .filter_map(|line| {
+      if line.is_empty() {
+        None
+      } else {
+        Some(
+          serde_json::from_slice::<PlaylistItem>(line.as_bytes()).map_err(
+            |err| ServiceError { status: 500, message: err.to_string() }
+          )
+        )
+      }
+    })
+    .filter_map(|parsed| match parsed {
+      Ok(item) => Some(item),
+      Err(err) => {
+        warn!("Failed to parse playlist item: {}", err);
+        None
+      }
+    })
+    .collect();
+  Ok(items)
 }
 
-fn is_valid_url(url: &str) -> bool {
-  match url.parse::<reqwest::Url>() {
-    Ok(_) => return true,
-    Err(_) => return false
-  }
+fn is_valid_url(url: &str) -> (bool, bool) {
+  Url::parse(url).ok().map_or((false, false), |valid_url| {
+    let is_playlist: bool = valid_url.query_pairs().find(|(key, _)| key == "list").map_or(false, |_| true);
+    (true, is_playlist)
+  })
 }
 
 pub async fn get_songbird(ctx: &Context) -> Arc<Songbird> {

service/src/bot/commands/audio/play.rs

@@ -9,9 +9,10 @@ use serenity::model::application::interaction::application_command::ApplicationC
 use siren::ServiceError;
 use songbird::{EventHandler, Songbird};
 
+use crate::bot::ytdlp::PlaylistItem;
 use crate::bot::{guilds::QueryGuild, commands::audio::{leave, get_playlist_urls, add_song, get_songbird}};
 
-use super::{create_response, edit_response, join_by_user};
+use super::{create_response, edit_response, is_valid_url, join_by_user};
 
 pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
   // Get the track url
@@ -65,8 +66,14 @@ pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
       debug!("Play command executed with track: {:?}", track_url);
       let manager = get_songbird(ctx).await;
       match play_track(manager, guild_id, track_url).await {
-        Ok(_) => {
-          if let Err(why) = edit_response(&ctx, &command, "Playing track".to_string()).await {
+        Ok(count) => {
+          let mut message = format!("Playing {} tracks", count);
+          if count == 0 {
+            message = "No tracks were played".to_string();
+          } else if count == 1 {
+            message = "Playing 1 track".to_string();
+          }
+          if let Err(why) = edit_response(&ctx, &command, message).await {
             error!("Failed to edit response message: {}", why);
           }
         },
@@ -95,15 +102,32 @@ pub async fn play_track(manager: Arc<Songbird>, guild_id: GuildId, track_url: St
       call_handler.queue().is_empty()
     };
     let guild = QueryGuild::get(guild_id.0 as i64)?;
-    let track_urls = match get_playlist_urls(&track_url) {
-      Ok(urls) => urls,
-      Err(err) => {
-        warn!("Failed to get playlist urls: {}", err);
-        return Err(ServiceError { status: 422, message: err.to_string() })
-      }
-    };
-    for url in track_urls {
-      match add_song(handler_lock.clone(), &url, is_queue_empty, Some(guild.volume as f32 / 100.0)).await {
+    let valid = is_valid_url(&track_url);
+    if !valid.0 {
+      warn!("Invalid track url: {}", track_url);
+      return Err(ServiceError { status: 422, message: format!("Invalid track url: {}", track_url) })
+    }
+    let mut playlist_items: Vec<PlaylistItem> = Vec::new();
+    if valid.1 {
+      playlist_items = match get_playlist_urls(&track_url) {
+        Ok(items) => items,
+        Err(err) => {
+          warn!("Failed to get playlist urls: {}", err);
+          return Err(ServiceError { status: 422, message: err.to_string() })
+        }
+      };
+    } else {
+      let playlist_item = PlaylistItem {
+        id: "".to_string(),
+        url: track_url,
+        title: "".to_string(),
+        duration: 0,
+        playlist_index: 0
+      };
+      playlist_items.push(playlist_item);
+    }
+    for item in playlist_items {
+      match add_song(handler_lock.clone(), &item.url, is_queue_empty, Some(guild.volume as f32 / 100.0)).await {
         Ok(added_song) => {
           let track_title = added_song.title.unwrap();
           debug!("Added track: {}", track_title);

service/src/bot/commands/chat.rs

@@ -1,4 +1,4 @@
-use log::{error, debug, warn};
+use log::{error, trace, warn};
 
 use serenity::model::Permissions;
 use serenity::model::channel::Message;
@@ -9,7 +9,7 @@ use crate::bot::messages::{QueryFilters, QueryMessage};
 use crate::bot::oai::{ChatCompletionMessage, ChatCompletionRequest, GPTRole, OAI};
 
 pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
-  debug!("Generating response for message: {}", msg.content);
+  trace!("Generating response for message: {}", msg.content);
 
   let guild_id = msg.guild_id.unwrap();
   let channel_id = msg.channel_id;
@@ -90,7 +90,7 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
   // Get the OAI response and store message/response into the database
   let response = match oai.chat_completion(request).await {
     Ok(r) => {
-      debug!("Processing response received from OpenAI");
+      trace!("Processing response received from OpenAI");
       if !r.choices.is_empty() {
         let res = r.choices[0].message.content.clone();
         if let Err(err) = QueryMessage::insert(QueryMessage {
@@ -118,7 +118,7 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
       "There was an error processing your message. Please try again later.".to_string()
     }
   };
-  debug!("Writing response: \"{}\"", response);
+  trace!("Writing response: \"{}\"", response);
 
   typing.stop();
   if let Err(why) = response_channel.say(&ctx.http, response).await {
@@ -142,7 +142,6 @@ pub async fn generate_response(ctx: &Context, msg: &Message, oai: &OAI) {
 }
 
 async fn generate_thread_name(oai: &OAI, s: &str, max_chars: usize) -> String {
-  println!("HERE: {}", s);
   let message = ChatCompletionMessage {
     role: GPTRole::User,
     content: format!("---\n{}\n---\nSummarize the message above into a concise Discord thread title", s)

service/src/bot/commands/mod.rs

@@ -3,3 +3,4 @@ pub mod help;
 pub mod chat;
 pub mod ping;
 pub mod schedule;
+pub mod roll;

service/src/bot/commands/roll.rs

@@ -0,0 +1,134 @@
+use log::{error, warn};
+use rand::Rng;
+use serenity::{builder::CreateApplicationCommand, client::Context, model::application::{command::CommandOptionType, interaction::application_command::ApplicationCommandInteraction}};
+
+use crate::bot::commands::audio::edit_response;
+
+use super::audio::create_response;
+
+pub async fn run(ctx: &Context, command: &ApplicationCommandInteraction) {
+  if let Err(why) = create_response(&ctx, &command, format!("Processing command...")).await {
+    error!("Failed to create response message: {}", why);
+    return;
+  }
+  let dice_string: String = match command.data.options.get(0) {
+    Some(o) => match &o.value {
+      Some(v) => match v.as_str() {
+        Some(s) => s.split_whitespace().collect::<String>(),
+        None => {
+          warn!("Missing dice option");
+          if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
+            error!("Failed to create response message: {}", why);
+          }
+          return;
+        }
+      },
+      None => {
+        warn!("Missing dice option");
+        if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
+          error!("Failed to create response message: {}", why);
+        }
+        return;
+      }
+    },
+    None => {
+      warn!("Missing dice option");
+      if let Err(why) = edit_response(&ctx, &command, format!("Dice option is missing")).await {
+        error!("Failed to create response message: {}", why);
+      }
+      return;
+    }
+  };
+  let dice = parse_dice(dice_string.as_str());
+  match dice {
+    Ok((count, sides, modifier)) => {
+      let mut rolls = Vec::new();
+      let mut total = 0;
+      for _ in 0..count {
+        let roll = rand::thread_rng().gen_range(1..=sides);
+        total += roll;
+        rolls.push(roll);
+      }
+      let response = format!("{}d{}{} = {}",
+        count,
+        sides, 
+        if modifier > 0 { format!("+{}", modifier) } else if modifier < 0 { format!("-{}", modifier) } else { "".to_string() },
+        total + (modifier as u32)
+      );
+      if let Err(why) = edit_response(&ctx, &command, response).await {
+        error!("Failed to create response message: {}", why);
+      }
+    }
+    Err(why) => {
+      if let Err(why) = edit_response(&ctx, &command, format!("Invalid dice string: {}", why)).await {
+        error!("Failed to create response message: {}", why);
+      }
+    }
+  }
+
+  
+}
+
+fn parse_dice(dice: &str) -> Result<(u32, u32, i32), String> {
+  let mut parts = dice.split("d");
+  let count = match parts.next() {
+    Some(c) => match c.parse::<u32>() {
+      Ok(n) => n,
+      Err(_) => return Err(format!("Invalid dice count: {}", c))
+    },
+    None => return Err(format!("Invalid dice string: {}", dice))
+  };
+  let mut positive_modifier = true;
+  let mut parts = match parts.next() {
+    Some(p) => {
+      // Check if contains a +/- modifier
+      if p.contains("+") {
+        positive_modifier = true;
+        p.split("+")
+      } else if p.contains("-") {
+        positive_modifier = false;
+        p.split("-")
+      } else {
+        p.split("+")
+      }
+    },
+    None => return Err(format!("Invalid dice string: {}", dice))
+  };
+  let sides = match parts.next() {
+    Some(s) => match s.parse::<u32>() {
+      Ok(n) => {
+        if n == 4 || n == 6 || n == 8 || n == 10 || n == 12 || n == 20 || n == 100 {
+          n
+        } else {
+          return Err(format!("Invalid dice sides: {}", s));
+        }
+      }
+      Err(_) => return Err(format!("Invalid dice sides: {}", s))
+    },
+    None => return Err(format!("Invalid dice string: {}", dice))
+  };
+  let modifier = match parts.next() {
+    Some(m) => match m.parse::<i32>() {
+      Ok(n) => {
+        if positive_modifier {
+          n
+        } else {
+          n * -1
+        }
+      },
+      Err(_) => return Err(format!("Invalid dice modifier: {}", m))
+    },
+    None => 0
+  };
+  Ok((count, sides, modifier))
+}
+
+pub fn register(command: &mut CreateApplicationCommand) -> &mut CreateApplicationCommand {
+  command.name("roll").description("Rolls D&D dice").create_option(|option| {
+    option
+      .name("dice")
+      .description("Dice to roll")
+      .kind(CommandOptionType::String)
+      .required(true)
+  })
+}