use crate::error::Error; use crate::error::Error::NoAlbumFound; use crate::imgur; use crate::imgur::Image; use crate::insult_compliment::InsultComplimentTemplate; use crate::wallet::WalletManager; use config::{Config, File}; use rand::prelude::SliceRandom; use serde::{Deserialize, Serialize}; use serenity::model::prelude::UserId; use serenity::prelude::TypeMapKey; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::Arc; use structopt::StructOpt; use tera::Tera; use tokio::sync::mpsc::{channel, Receiver, Sender}; use tokio::sync::Mutex; #[derive(Debug, StructOpt)] #[structopt(name = "fren", about = "Friend Bot")] pub struct Args { pub cfg_path: PathBuf, } #[derive(Debug, Deserialize, Serialize, Clone)] pub struct AlbumConfig { pub name: String, pub album_id: String, } #[derive(Debug, Deserialize, Serialize, Clone)] pub struct RandomConfig { pub name: String, pub responses: Vec, } #[derive(Debug, Deserialize, Serialize, Clone)] pub struct BotConfig { pub bot_token: String, pub imgur_client_id: String, pub story_path: PathBuf, #[serde(default)] pub albums: Vec, #[serde(default)] pub randoms: Vec, pub wallet_manager: WalletManager, pub insults: Vec, pub compliments: Vec, } impl BotConfig { pub fn new(config_path: &Path) -> Result { let cfg = Config::builder() .add_source(File::from(config_path)) .build()?; cfg.try_deserialize() } pub async fn save(&self, config_path: &Path) -> Result<(), tokio::io::Error> { let output_str = toml::to_string_pretty(&self).unwrap(); tokio::fs::write(config_path, output_str).await?; Ok(()) } } #[derive(Debug)] pub struct BotState { pub accepted_nsfw: Option, pub templates: HashMap, pub albums: HashMap>, pub bad_apple_running: bool, pub story_in_progress: Mutex<()>, } impl BotState { pub async fn new(cfg: &BotConfig) -> Result { let mut albums: HashMap> = HashMap::new(); let mut templates = HashMap::new(); for random_config in cfg.randoms.iter() { let mut templates_set = Tera::default(); for (idx, response) in random_config.responses.iter().enumerate() { templates_set .add_raw_template(&idx.to_string(), response) .unwrap(); } templates.insert(random_config.name.clone(), templates_set); } for album in &cfg.albums { albums.insert( album.name.clone(), imgur::get_album_images(&cfg.imgur_client_id, &album.album_id).await?, ); } Ok(Self { accepted_nsfw: None, templates, albums, bad_apple_running: false, story_in_progress: 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(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()) } } #[derive(Debug)] pub struct GlobalData { pub args: Args, pub cfg: BotConfig, pub bot_state: BotState, } impl GlobalData { pub async fn new(args: Args, cfg: BotConfig) -> Result { Ok(Self { args, bot_state: BotState::new(&cfg).await?, cfg, }) } pub async fn reload(&mut self) -> Result<(), Error> { let cfg = BotConfig::new(&self.args.cfg_path)?; self.cfg = cfg; self.bot_state = BotState::new(&self.cfg).await?; Ok(()) } } impl TypeMapKey for GlobalData { type Value = GlobalData; } #[derive(Debug, Clone)] pub struct Channel { pub recv: Arc>>, pub send: Arc>>, } impl Channel { pub fn new() -> Self { let (send, recv) = channel::(10); Self { recv: Arc::new(Mutex::new(recv)), send: Arc::new(Mutex::new(send)), } } } impl TypeMapKey for Channel { type Value = Channel; }