FrenBot/src/config.rs

use crate::error::Error;
use crate::error::Error::NoAlbumFound;
use crate::imgur;
use crate::imgur::Image;
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<String>,
}

#[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<AlbumConfig>,

    #[serde(default)]
    pub randoms: Vec<RandomConfig>,

    pub wallet_manager: WalletManager,
}

impl BotConfig {
    pub fn new(config_path: &Path) -> Result<Self, config::ConfigError> {
        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<UserId>,
    pub templates: HashMap<String, Tera>,
    pub albums: HashMap<String, Vec<Image>>,
    pub bad_apple_running: bool,
    pub story_in_progress: Mutex<()>,
}

impl BotState {
    pub async fn new(cfg: &BotConfig) -> Result<Self, Error> {
        let mut albums: HashMap<String, Vec<Image>> = 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<Option<Image>, 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<Self, Error> {
        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<T> {
    pub recv: Arc<Mutex<Receiver<T>>>,
    pub send: Arc<Mutex<Sender<T>>>,
}

impl<T> Channel<T> {
    pub fn new() -> Self {
        let (send, recv) = channel::<T>(10);

        Self {
            recv: Arc::new(Mutex::new(recv)),
            send: Arc::new(Mutex::new(send)),
        }
    }
}

impl<T: 'static + Send + Sync> TypeMapKey for Channel<T> {
    type Value = Channel<T>;
}