FrenBot/src/config.rs

use crate::album_manager::AlbumManager;
use crate::error::Error;
use crate::event_listener::{Action, Expiration, Listener, TriggerType};
use crate::migrations::{CURRENT_DB_VERSION, do_migration};
use config::{Config, File};
use cta_api::CTAClient;
use j_db::database::{DB_METADATA_ID, Database};
use j_db::metadata::DBMetadata;
use j_db::model::JdbModel;
use log::info;
use poise::serenity_prelude::model::id::ChannelId;
use poise::serenity_prelude::model::prelude::{GuildId, UserId};
use poise::serenity_prelude::prelude::TypeMapKey;
use reqwest::Url;
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use structopt::StructOpt;
use tokio::sync::Mutex;

#[derive(Debug, StructOpt)]
#[structopt(name = "fren", about = "Friend Bot")]
pub struct Args {
    pub cfg_path: PathBuf,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PicOxConfig {
    api_base_url: Url,
    token: String,
}

#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct BotConfig {
    pub bot_token: String,
    pub img_path: PathBuf,
    pub base_url: Url,
    pub story_path: PathBuf,
    pub voice_path: PathBuf,
    pub nft_path: PathBuf,
    pub db_path: PathBuf,
    pub admins: Vec<UserId>,
    pub guild_id: GuildId,
    pub api_addr: SocketAddr,
    pub announcement_channel: ChannelId,
    pub toys: Vec<String>,
    pub effect_role_duration: i64,
    pub raas_server: String,
    pub cta_key: String,
    pub omdb_key: String,
    pub music_streaming_site_urls: Vec<String>,
    pub music_streaming_sites: Vec<String>,
    pub picox: PicOxConfig,
}

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 fn effect_role_duration(&self) -> chrono::Duration {
        chrono::Duration::seconds(self.effect_role_duration)
    }
}

#[derive(Debug)]
pub struct BotState {
    pub accepted_nsfw: Option<UserId>,
}

impl BotState {
    pub async fn new() -> Result<Self, Error> {
        Ok(Self {
            accepted_nsfw: None,
        })
    }
}

pub struct GlobalData {
    #[allow(dead_code)]
    pub args: Args,
    pub cfg: BotConfig,
    pub bot_state: Mutex<BotState>,
    pub db: Database,
    pub picox: AlbumManager,
    pub cta: Mutex<CTAClient>,
    pub speak_lock: Mutex<()>,
    pub listener_lock: Mutex<()>,
}

impl GlobalData {
    fn setup_system_listeners(db: &Database) -> Result<(), Error> {
        db.filter(|_, listener: &Listener| listener.system_listener)?
            .for_each(|listener: Listener| {
                let _ = db.remove::<Listener>(listener.id().unwrap()).is_ok();
            });

        info!("Adding system listeners...");
        Listener::add_listener(
            db,
            Listener::new(
                TriggerType::OnMessage {
                    channel_id: None,
                    content: Some("bad bot".to_string()),
                },
                vec![Action::React {
                    emoji: "😭".to_string(),
                }],
                1.0,
                Expiration::Never,
                true,
            ),
        )?;
        Listener::add_listener(
            db,
            Listener::new(
                TriggerType::OnMessage {
                    channel_id: None,
                    content: None,
                },
                vec![Action::UpdateFrenCoins {
                    fren_coin_diff: 100,
                }],
                0.05,
                Expiration::Never,
                true,
            ),
        )?;

        Ok(())
    }

    pub async fn new(args: Args, cfg: BotConfig) -> Result<Self, Error> {
        let db = Database::new(&cfg.db_path)?;

        let version = db.version()?;

        if version == 0 {
            db.insert::<j_db::metadata::DBMetadata>(DBMetadata {
                id: Some(DB_METADATA_ID),
                version: CURRENT_DB_VERSION,
            })?;
        }

        do_migration(&db);

        Self::setup_system_listeners(&db)?;

        Ok(Self {
            args,
            bot_state: Mutex::new(BotState::new().await?),
            db,
            cfg: cfg.clone(),
            picox: AlbumManager::new(cfg.picox.api_base_url, &cfg.picox.token),
            cta: Mutex::new(CTAClient::new(cfg.cta_key)),
            speak_lock: Default::default(),
            listener_lock: Default::default(),
        })
    }
}

impl TypeMapKey for GlobalData {
    type Value = GlobalData;
}