diff --git a/src/discord/motivate.rs b/src/discord/motivate.rs index 9f71680..128710c 100644 --- a/src/discord/motivate.rs +++ b/src/discord/motivate.rs @@ -59,13 +59,23 @@ pub async fn create_motivation_image(motivation: Motivation) -> Result, #[command] #[description("Let's give you motivation")] -async fn motivation(ctx: &Context, msg: &Message, _args: Args) -> CommandResult { +async fn motivation(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let data = ctx.data.read().await; let global_data = data.get::().unwrap(); - let motivation = - MotivationConfig::generate_motivation(&global_data.db, &global_data.cfg.img_path, "white") - .await?; + let album_name = if args.is_empty() { + None + } else { + Some(args.parse::()?) + }; + + let motivation = MotivationConfig::generate_motivation( + &global_data.db, + &global_data.cfg.img_path, + "white", + album_name, + ) + .await?; let image = create_motivation_image(motivation).await?; diff --git a/src/discord/shop.rs b/src/discord/shop.rs index 17fcca4..6e94eef 100644 --- a/src/discord/shop.rs +++ b/src/discord/shop.rs @@ -322,6 +322,7 @@ pub async fn restock_shop(ctx: &Context) -> Result<(), CommandError> { &global_data.db, &global_data.cfg.img_path, "gold", + None, ) .await .unwrap(); diff --git a/src/models/motivation.rs b/src/models/motivation.rs index e4eeb1d..d39925d 100644 --- a/src/models/motivation.rs +++ b/src/models/motivation.rs @@ -1,6 +1,7 @@ use crate::album_manager::Album; use crate::error::Error; use j_db::database::Database; +use j_db::error::JDbError; use j_db::model::JdbModel; use rand::prelude::{IteratorRandom, SliceRandom}; use rand::thread_rng; @@ -80,17 +81,28 @@ impl MotivationConfig { db: &Database, base_path: &Path, border_color: &str, + album_name: Option, ) -> Result { let motivation = Self::get_motivation_config(db)?; let mut images = Vec::new(); - for album in motivation.album { - let album = db.get::(album); + if let Some(album_name) = album_name { + let album = Album::find_album_by_name_or_alias(db, &album_name)?; - if let Ok(album) = album { - let mut album_images = album.images.clone(); - images.append(&mut album_images); + if let Some(mut album) = album { + images.append(&mut album.images) + } else { + return Err(Error::from(JDbError::NotFound)); + } + } else { + for album in motivation.album { + let album = db.get::(album); + + if let Ok(album) = album { + let mut album_images = album.images.clone(); + images.append(&mut album_images); + } } }