From 7847bd38b95d5d0fc336d6bb4c4159ec85d5a47b Mon Sep 17 00:00:00 2001 From: Joey Hines Date: Sun, 4 Jun 2023 14:10:01 -0600 Subject: [PATCH] Removed imgur support + Added my own basic image handling + Probably not perfect but works good enough in testing + clippy + fmt --- src/album_manager/mod.rs | 126 ++++++++++++++++++ src/config.rs | 67 ++-------- src/discord/album.rs | 151 ++++++++++++++-------- src/discord/joke.rs | 11 +- src/discord/motivate.rs | 11 +- src/discord/shop.rs | 2 +- src/error.rs | 9 -- src/imgur/mod.rs | 75 ----------- src/main.rs | 5 +- src/migrations/migration2_remove_imgur.rs | 43 ++++++ src/migrations/mod.rs | 26 ++++ src/models/album.rs | 41 ------ src/models/mod.rs | 1 - src/models/motivation.rs | 28 ++-- 14 files changed, 339 insertions(+), 257 deletions(-) create mode 100644 src/album_manager/mod.rs delete mode 100644 src/imgur/mod.rs create mode 100644 src/migrations/migration2_remove_imgur.rs create mode 100644 src/migrations/mod.rs delete mode 100644 src/models/album.rs diff --git a/src/album_manager/mod.rs b/src/album_manager/mod.rs new file mode 100644 index 0000000..a7a12c4 --- /dev/null +++ b/src/album_manager/mod.rs @@ -0,0 +1,126 @@ +use crate::error::Error; +use j_db::database::Database; +use j_db::error::JDbError; +use j_db::model::JdbModel; +use rand::prelude::SliceRandom; +use reqwest::Url; +use serde::{Deserialize, Serialize}; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::path::PathBuf; + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Image { + pub filename: String, + pub tags: Vec, + pub hash: u64, + pub db_name: String, +} + +impl Image { + pub async fn new(data: &[u8], filename: &str, base_dir: PathBuf, tags: Vec) -> Self { + let mut hasher = DefaultHasher::new(); + data.hash(&mut hasher); + filename.hash(&mut hasher); + let hash = hasher.finish(); + let file = PathBuf::from(filename); + let ext = file.extension().unwrap().to_str().unwrap(); + + let db_name = format!("{}.{}", hash, ext); + let path = base_dir.join(db_name.clone()); + tokio::fs::write(path, data).await.unwrap(); + + Self { + filename: filename.to_string(), + tags, + hash, + db_name, + } + } + + pub fn full_path(&self, base_dir: PathBuf) -> PathBuf { + base_dir.join(&self.db_name) + } + + pub fn link(&self, base_link: &Url) -> Url { + base_link.join(&self.db_name).unwrap() + } +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Album { + id: Option, + pub album_name: String, + pub images: Vec, + pub aliases: Vec, +} + +impl Album { + pub fn find_album_by_name_or_alias( + db: &Database, + name: &str, + ) -> Result, JDbError> { + Ok(db + .filter(|_, album: &Album| { + album.album_name.eq_ignore_ascii_case(name) + || album.aliases.contains(&name.to_string()) + })? + .next()) + } + + pub fn get_random_image( + db: &Database, + album_name: &str, + tags: Vec<&str>, + ) -> Result, Error> { + let mut rng = rand::thread_rng(); + let album = match Self::find_album_by_name_or_alias(db, album_name)? { + None => return Err(Error::NoAlbumFound), + Some(a) => a, + }; + + let images: Vec = if tags.is_empty() { + album.images + } else { + album + .images + .iter() + .filter(|img| { + for tag in &tags { + if img.tags.contains(&tag.to_lowercase()) { + return true; + } + } + + false + }) + .cloned() + .collect() + }; + + Ok(images.choose(&mut rng).cloned()) + } + + pub fn add_album(db: &Database, album_name: &str) -> Result { + db.insert(Album { + id: None, + album_name: album_name.to_string(), + images: vec![], + aliases: vec![], + }) + } +} + +impl JdbModel for Album { + fn id(&self) -> Option { + self.id + } + + fn set_id(&mut self, id: u64) { + self.id = Some(id) + } + + fn tree() -> String { + "img_album".to_string() + } +} diff --git a/src/config.rs b/src/config.rs index 3106ede..79bffb1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,15 +1,12 @@ use crate::error::Error; -use crate::imgur; -use crate::imgur::Image; -use crate::models::album::AlbumConfig; +use crate::migrations::do_migration; use config::{Config, File}; use j_db::database::Database; -use rand::prelude::SliceRandom; +use reqwest::Url; use serde::{Deserialize, Serialize}; use serenity::model::id::ChannelId; use serenity::model::prelude::{GuildId, UserId}; use serenity::prelude::TypeMapKey; -use std::collections::HashMap; use std::net::SocketAddr; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -26,7 +23,8 @@ pub struct Args { #[derive(Debug, Deserialize, Serialize, Clone)] pub struct BotConfig { pub bot_token: String, - pub imgur_client_id: String, + pub img_path: PathBuf, + pub base_url: Url, pub story_path: PathBuf, pub voice_path: PathBuf, pub nft_path: PathBuf, @@ -50,66 +48,18 @@ impl BotConfig { #[derive(Debug)] pub struct BotState { pub accepted_nsfw: Option, - pub albums: HashMap>, pub bad_apple_running: bool, pub speak_lock: Mutex<()>, } impl BotState { - pub async fn new(cfg: &BotConfig, db: &Database) -> Result { - let mut albums: HashMap> = HashMap::new(); - - for album in db.filter(|_, _: &AlbumConfig| true)? { - let imgur_album = imgur::get_album_images(&cfg.imgur_client_id, &album.album_id).await; - - match imgur_album { - Ok(imgur) => { - albums.insert(album.name.clone(), imgur); - } - Err(err) => { - println!("Unable to parse album '{}'. Error: {}", album.name, err); - } - } - } - + pub async fn new() -> Result { Ok(Self { accepted_nsfw: None, - albums, bad_apple_running: false, speak_lock: Mutex::new(()), }) } - - pub fn get_image(&self, album_name: &str, tags: Vec<&str>) -> Result, Error> { - let mut rng = rand::thread_rng(); - - let album = match self.albums.get(album_name) { - None => return Err(Error::NoAlbumFound), - Some(a) => a, - }; - - let album = if tags.is_empty() { - album.clone() - } else { - album - .iter() - .filter(|img| { - for tag in &tags { - if let Some(desc) = &img.description { - if desc.to_lowercase().contains(&tag.to_lowercase()) { - return true; - } - } - } - - false - }) - .cloned() - .collect() - }; - - Ok(album.choose(&mut rng).cloned()) - } } pub struct GlobalData { @@ -122,9 +72,12 @@ pub struct GlobalData { impl GlobalData { pub async fn new(args: Args, cfg: BotConfig) -> Result { let db = Database::new(&cfg.db_path)?; + + do_migration(&db); + Ok(Self { args, - bot_state: BotState::new(&cfg, &db).await?, + bot_state: BotState::new().await?, db, cfg, }) @@ -134,7 +87,7 @@ impl GlobalData { let cfg = BotConfig::new(&self.args.cfg_path)?; self.cfg = cfg; - self.bot_state = BotState::new(&self.cfg, &self.db).await?; + self.bot_state = BotState::new().await?; Ok(()) } diff --git a/src/discord/album.rs b/src/discord/album.rs index 2492dfe..9c5e259 100644 --- a/src/discord/album.rs +++ b/src/discord/album.rs @@ -1,6 +1,6 @@ +use crate::album_manager::{Album, Image}; +use crate::discord::admin::is_admin; use crate::error::Error; -use crate::imgur::get_album_images; -use crate::models::album::AlbumConfig; use crate::{command, group, GlobalData}; use j_db::model::JdbModel; use serenity::client::Context; @@ -8,64 +8,106 @@ use serenity::framework::standard::{Args, CommandResult}; use serenity::model::channel::Message; #[group] -#[commands(add_album, remove_album, list_albums)] -pub struct Album; +#[commands(add_image, remove_album, list_albums, import_from_file)] +pub struct AlbumCmd; #[command] #[only_in(guilds)] -#[min_args(2)] -#[max_args(2)] -#[description("Add an imgur album command.")] -#[usage(" ")] -async fn add_album(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { +#[min_args(1)] +#[description("Add an image to an album")] +#[usage("")] +async fn add_image(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { let album_name = args.parse::()?; args.advance(); - let album_id = args.parse::()?; - - let album_id: String = if album_id.contains("imgur") { - let parts: Vec<&str> = album_id.split("/a/").collect(); - - if parts.len() == 2 { - parts[1].to_string() - } else { - msg.reply(&ctx.http, "Invalid imgur album, check your link.") - .await?; - return Ok(()); - } - } else { - album_id - }; + let tags: Vec = args.rest().split(',').map(|s| s.to_string()).collect(); let mut data = ctx.data.write().await; let global_data = data.get_mut::().unwrap(); - let images = match get_album_images(&global_data.cfg.imgur_client_id, &album_id).await { - Ok(img) => img, - Err(e) => { - msg.reply( - &ctx.http, - "Error adding album, check your link and try again", - ) - .await?; + let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap(); - println!("Error adding album: {}", e); - - return Ok(()); - } + let mut album = if let Some(album) = album { + album + } else { + Album::add_album(&global_data.db, &album_name).unwrap() }; - global_data - .db - .insert(AlbumConfig::new(&album_name, &album_id))?; + for attachment in &msg.attachments { + let data = attachment.download().await.unwrap(); + let img = Image::new( + &data, + &attachment.filename, + global_data.cfg.img_path.clone(), + tags.clone(), + ) + .await; + album.images.push(img); + } - global_data - .bot_state - .albums - .insert(album_name.clone(), images); + global_data.db.insert(album).unwrap(); - msg.reply(&ctx.http, format!("{} album added!", album_name)) - .await?; + let plural = if msg.attachments.len() > 1 { "s" } else { "" }; + + msg.reply( + &ctx.http, + format!("Image{} added to {}!", plural, album_name), + ) + .await?; + + Ok(()) +} + +#[command] +#[only_in(guilds)] +#[min_args(1)] +#[description("Bulk add images")] +#[usage("")] +async fn import_from_file(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { + let album_name = args.parse::()?; + args.advance(); + let album_path = args.parse::()?; + let album_path = album_path.replace('"', ""); + + let mut data = ctx.data.write().await; + + let global_data = data.get_mut::().unwrap(); + + if !is_admin(&msg.author.id, &global_data.cfg) { + return Ok(()); + } + + let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap(); + + let mut album = if let Some(album) = album { + album + } else { + Album::add_album(&global_data.db, &album_name).unwrap() + }; + + let mut dir = tokio::fs::read_dir(album_path).await.unwrap(); + + while let Some(entry) = dir.next_entry().await? { + let data = tokio::fs::read(entry.path()).await?; + let img = Image::new( + &data, + entry.path().file_name().unwrap().to_str().unwrap(), + global_data.cfg.img_path.clone(), + vec![], + ) + .await; + album.images.push(img); + } + + global_data.db.insert(album).unwrap(); + + let plural = if msg.attachments.len() > 1 { "s" } else { "" }; + + msg.reply( + &ctx.http, + format!("Image{} added to {}!", plural, album_name), + ) + .await?; Ok(()) } @@ -82,7 +124,7 @@ async fn remove_album(ctx: &Context, msg: &Message, args: Args) -> CommandResult let global_data = data.get_mut::().unwrap(); - let album = AlbumConfig::get_album_by_name(&global_data.db, &album_name)?; + let album = Album::find_album_by_name_or_alias(&global_data.db, &album_name).unwrap(); match album { None => { @@ -93,9 +135,13 @@ async fn remove_album(ctx: &Context, msg: &Message, args: Args) -> CommandResult .await?; } Some(album) => { - global_data.bot_state.albums.remove(&album_name); + for img in &album.images { + tokio::fs::remove_file(img.full_path(global_data.cfg.img_path.clone())) + .await + .unwrap(); + } - global_data.db.remove::(album.id().unwrap())?; + global_data.db.remove::(album.id().unwrap())?; msg.reply(&ctx.http, format!("{} album removed!", album_name)) .await?; @@ -116,8 +162,8 @@ async fn list_albums(ctx: &Context, msg: &Message, _args: Args) -> CommandResult let album_names: Vec = global_data .db - .filter(|_, _: &AlbumConfig| true)? - .map(|album| album.name) + .filter(|_, _: &Album| true)? + .map(|album| album.album_name) .collect(); if album_names.is_empty() { @@ -143,7 +189,7 @@ pub async fn parse_album( let data = ctx.data.read().await; let global_data = data.get::().unwrap(); - let img = match global_data.bot_state.get_image(album_name, tags) { + let img = match Album::get_random_image(&global_data.db, album_name, tags) { Ok(img) => img, Err(err) => { return match err { @@ -154,7 +200,8 @@ pub async fn parse_album( }; if let Some(img) = img { - msg.reply(&ctx.http, img.link).await?; + msg.reply(&ctx.http, img.link(&global_data.cfg.base_url)) + .await?; } else { msg.reply(&ctx.http, "No image found :(").await?; } diff --git a/src/discord/joke.rs b/src/discord/joke.rs index b4bd0e7..db7cc75 100644 --- a/src/discord/joke.rs +++ b/src/discord/joke.rs @@ -1,3 +1,4 @@ +use crate::album_manager::Album; use crate::error::Error; use crate::models::insult_compliment::{RandomResponseTemplate, ResponseType}; use crate::models::random::RandomConfig; @@ -55,11 +56,15 @@ impl RandomCtx { pub fn new(user_name: &str, global_data: &GlobalData) -> Result { let mut random_image: HashMap = HashMap::new(); - for (album_name, images) in &global_data.bot_state.albums { - let image = images.choose(&mut thread_rng()); + let albums: Vec = global_data.db.filter(|_, _album: &Album| true)?.collect(); + for album in albums { + let image = album.images.choose(&mut thread_rng()); if let Some(image) = image { - random_image.insert(album_name.clone(), image.link.clone()); + random_image.insert( + album.album_name.clone(), + image.link(&global_data.cfg.base_url).to_string(), + ); } } diff --git a/src/discord/motivate.rs b/src/discord/motivate.rs index 7187221..b829a73 100644 --- a/src/discord/motivate.rs +++ b/src/discord/motivate.rs @@ -16,7 +16,7 @@ use std::borrow::Cow; pub struct Motivate; pub async fn create_motivation_image(motivation: Motivation) -> Result, CommandError> { - let motivation_image_blob = reqwest::get(&motivation.image_url).await?.bytes().await?; + let motivation_image_blob = tokio::fs::read(motivation.image_path).await?; let text = format!("{} {}", motivation.action, motivation.goal); @@ -64,12 +64,9 @@ 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.imgur_client_id, - "white", - ) - .await?; + let motivation = + MotivationConfig::generate_motivation(&global_data.db, &global_data.cfg.img_path, "white") + .await?; let image = create_motivation_image(motivation).await?; diff --git a/src/discord/shop.rs b/src/discord/shop.rs index a3824d9..f0bc6c6 100644 --- a/src/discord/shop.rs +++ b/src/discord/shop.rs @@ -312,7 +312,7 @@ pub async fn restock_shop(ctx: &Context) -> Result<(), CommandError> { let nft_motivation = MotivationConfig::generate_motivation( &global_data.db, - &global_data.cfg.imgur_client_id, + &global_data.cfg.img_path, "gold", ) .await diff --git a/src/error.rs b/src/error.rs index 70b21a4..a472544 100644 --- a/src/error.rs +++ b/src/error.rs @@ -1,4 +1,3 @@ -use crate::imgur::ImgurError; use crate::user; use serde::ser::StdError; use std::fmt::{Display, Formatter}; @@ -8,7 +7,6 @@ use std::fmt::{Display, Formatter}; #[allow(dead_code)] pub enum Error { ConfigError(config::ConfigError), - ImgurError(ImgurError), SerenityError(serenity::Error), TeraError(tera::Error), NoAlbumFound, @@ -24,12 +22,6 @@ impl From for Error { } } -impl From for Error { - fn from(e: ImgurError) -> Self { - Self::ImgurError(e) - } -} - impl From for Error { fn from(err: serenity::Error) -> Self { Self::SerenityError(err) @@ -58,7 +50,6 @@ impl Display for Error { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { Error::ConfigError(e) => write!(f, "Config error: {}", e), - Error::ImgurError(e) => write!(f, "Imgur error: {}", e), Error::SerenityError(e) => write!(f, "Discord error: {}", e), Error::TeraError(e) => write!(f, "Tera error: {}", e), Error::NoAlbumFound => write!(f, "No album found"), diff --git a/src/imgur/mod.rs b/src/imgur/mod.rs deleted file mode 100644 index 6f98f71..0000000 --- a/src/imgur/mod.rs +++ /dev/null @@ -1,75 +0,0 @@ -use reqwest::Client; -use serde::{Deserialize, Serialize}; -use std::fmt::{Display, Formatter}; - -#[derive(Debug)] -pub enum ImgurError { - ReqwestError(reqwest::Error), - ImgurRequestError(String), -} - -impl From for ImgurError { - fn from(e: reqwest::Error) -> Self { - Self::ReqwestError(e) - } -} - -impl Display for ImgurError { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - let msg = match self { - ImgurError::ReqwestError(err) => format!("Reqwest error: {}", err), - ImgurError::ImgurRequestError(msg) => format!("Imgur request error: {}", msg), - }; - - write!(f, "{}", msg) - } -} - -impl std::error::Error for ImgurError {} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct AlbumData { - images: Option>, - error: Option, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct AlbumResponse { - data: AlbumData, - success: bool, - status: i32, -} - -#[derive(Serialize, Deserialize, Debug, Clone)] -pub struct Image { - pub id: String, - pub title: Option, - pub description: Option, - #[serde(rename = "type")] - pub img_type: String, - pub animated: bool, - pub width: i32, - pub height: i32, - pub size: i32, - pub link: String, -} - -pub async fn get_album_images(client_id: &str, album_hash: &str) -> Result, ImgurError> { - let client = Client::new(); - - let res = client - .get(format!("https://api.imgur.com/3/album/{}", album_hash)) - .header("Authorization", format!("Client-ID {}", client_id)) - .send() - .await?; - - let album_response: AlbumResponse = res.json().await?; - - if album_response.success { - Ok(album_response.data.images.unwrap()) - } else { - Err(ImgurError::ImgurRequestError( - album_response.data.error.unwrap(), - )) - } -} diff --git a/src/main.rs b/src/main.rs index 2c592aa..2f15d24 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,10 @@ +mod album_manager; mod api; mod config; mod discord; mod error; -mod imgur; mod inventory; +mod migrations; mod models; mod user; @@ -49,7 +50,7 @@ async fn main() { let framework = StandardFramework::new() .configure(|c| c.with_whitespace(true).prefix("!").ignore_bots(true)) .group(&discord::color::COLOR_GROUP) - .group(&discord::album::ALBUM_GROUP) + .group(&discord::album::ALBUMCMD_GROUP) .group(&discord::celeryman::CELERYMAN_GROUP) .group(&discord::joke::JOKE_GROUP) .group(&discord::admin::ADMIN_GROUP) diff --git a/src/migrations/migration2_remove_imgur.rs b/src/migrations/migration2_remove_imgur.rs new file mode 100644 index 0000000..9cdfa46 --- /dev/null +++ b/src/migrations/migration2_remove_imgur.rs @@ -0,0 +1,43 @@ +use crate::models::motivation::MotivationConfig; +use j_db::database::Database; +use j_db::migration::Migration; +use j_db::model::JdbModel; +use json::JsonValue; + +pub struct Migration2RemoveImgur {} + +impl Migration for Migration2RemoveImgur { + fn up(&self, db: &Database) -> j_db::error::Result<()> { + db.db.open_tree("Album").unwrap().clear()?; + + let mot_tree = db.db.open_tree(MotivationConfig::tree())?; + + let (id, config) = mot_tree.iter().next().unwrap()?; + + let mut config = json::parse(std::str::from_utf8(&config).unwrap()).unwrap(); + + config["album"] = JsonValue::Array(Vec::new()); + + mot_tree.insert(id, config.to_string().as_bytes())?; + + Ok(()) + } + + fn down(&self, db: &Database) -> j_db::error::Result<()> { + let mot_tree = db.db.open_tree(MotivationConfig::tree())?; + + let (id, config) = mot_tree.iter().next().unwrap()?; + + let mut config = json::parse(std::str::from_utf8(&config).unwrap()).unwrap(); + + config["album"] = JsonValue::Array(Vec::new()); + + mot_tree.insert(id, config.to_string().as_bytes())?; + + Ok(()) + } + + fn version(&self) -> u64 { + 2 + } +} diff --git a/src/migrations/mod.rs b/src/migrations/mod.rs new file mode 100644 index 0000000..bb59e25 --- /dev/null +++ b/src/migrations/mod.rs @@ -0,0 +1,26 @@ +use crate::migrations::migration2_remove_imgur::Migration2RemoveImgur; +use j_db::database::Database; +use j_db::migration; +use j_db::migration::Direction; + +mod migration2_remove_imgur; + +const CURRENT_DB_VERSION: u64 = 2; + +#[allow(clippy::single_match)] +pub fn do_migration(db: &Database) { + let version = db.version().unwrap_or(0); + for ver in version + 1..=CURRENT_DB_VERSION { + match ver { + 2 => { + migration::do_migration::( + db, + Migration2RemoveImgur {}, + Direction::Up, + ) + .unwrap(); + } + _ => {} + } + } +} diff --git a/src/models/album.rs b/src/models/album.rs deleted file mode 100644 index 60b6ec6..0000000 --- a/src/models/album.rs +++ /dev/null @@ -1,41 +0,0 @@ -use crate::error::Error; -use j_db::database::Database; -use j_db::model::JdbModel; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Deserialize, Serialize, Clone)] -pub struct AlbumConfig { - id: Option, - pub name: String, - pub album_id: String, -} - -impl JdbModel for AlbumConfig { - fn id(&self) -> Option { - self.id - } - - fn set_id(&mut self, id: u64) { - self.id = Some(id) - } - - fn tree() -> String { - "Album".to_string() - } -} - -impl AlbumConfig { - pub fn new(name: &str, album_id: &str) -> AlbumConfig { - Self { - id: None, - name: name.to_string(), - album_id: album_id.to_string(), - } - } - - pub fn get_album_by_name(db: &Database, name: &str) -> Result, Error> { - Ok(db - .filter(|_, album_config: &AlbumConfig| album_config.name.eq_ignore_ascii_case(name))? - .next()) - } -} diff --git a/src/models/mod.rs b/src/models/mod.rs index 4d8f094..69c0871 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,4 +1,3 @@ -pub mod album; pub mod api_key; pub mod birthday; pub mod insult_compliment; diff --git a/src/models/motivation.rs b/src/models/motivation.rs index 12166de..e4eeb1d 100644 --- a/src/models/motivation.rs +++ b/src/models/motivation.rs @@ -1,15 +1,16 @@ +use crate::album_manager::Album; use crate::error::Error; -use crate::imgur::get_album_images; use j_db::database::Database; use j_db::model::JdbModel; use rand::prelude::{IteratorRandom, SliceRandom}; use rand::thread_rng; use serde::{Deserialize, Serialize}; +use std::path::{Path, PathBuf}; #[derive(Debug, Deserialize, Serialize, Clone)] pub struct MotivationConfig { id: Option, - pub album: Vec, + pub album: Vec, pub action: Vec, pub goal: Vec, } @@ -47,7 +48,10 @@ impl MotivationConfig { pub fn add_album(db: &Database, album: &str) -> Result<(), Error> { let mut motivation = Self::get_motivation_config(db)?; - motivation.album.push(album.to_string()); + let album = Album::find_album_by_name_or_alias(db, album) + .unwrap() + .unwrap(); + motivation.album.push(album.id().unwrap()); db.insert(motivation)?; @@ -74,7 +78,7 @@ impl MotivationConfig { pub async fn generate_motivation( db: &Database, - imgur_client_id: &str, + base_path: &Path, border_color: &str, ) -> Result { let motivation = Self::get_motivation_config(db)?; @@ -82,18 +86,24 @@ impl MotivationConfig { let mut images = Vec::new(); for album in motivation.album { - let mut album_images = get_album_images(imgur_client_id, &album).await?; - images.append(&mut album_images); + let album = db.get::(album); + + if let Ok(album) = album { + let mut album_images = album.images.clone(); + images.append(&mut album_images); + } } - let image = images.choose(&mut thread_rng()).unwrap(); + let image = images + .choose(&mut thread_rng()) + .ok_or(Error::NoAlbumFound)?; let action = motivation.action.iter().choose(&mut thread_rng()).unwrap(); let goal = motivation.goal.iter().choose(&mut thread_rng()).unwrap(); Ok(Motivation { - image_url: image.link.clone(), + image_path: image.full_path(base_path.to_path_buf()), action: action.to_string(), goal: goal.to_string(), border_color: border_color.to_string(), @@ -103,7 +113,7 @@ impl MotivationConfig { #[derive(Debug, Deserialize, Serialize, Clone)] pub struct Motivation { - pub image_url: String, + pub image_path: PathBuf, pub action: String, pub goal: String, pub border_color: String,