use std::{collections::HashMap, sync::Arc}; use async_trait::async_trait; use futures_util::{ future::{join_all, ready, FutureExt}, stream::{FuturesUnordered, StreamExt}, }; use serenity::{ client::Cache, framework::standard::{ macros::{command, group}, Args, CommandResult, }, http::{CacheHttp, Http}, model::{ channel::Message, id::{ChannelId, GuildId, UserId}, }, prelude::*, utils::MessageBuilder, }; use youmubot_db::DB; use crate::{AppData, MemberCache, Result}; #[derive(Debug, Clone)] pub struct CacheAndHttp(Arc, Arc); impl CacheAndHttp { pub fn from_client(client: &Client) -> Self { Self(client.cache.clone(), client.http.clone()) } pub fn from_context(context: &Context) -> Self { Self(context.cache.clone(), context.http.clone()) } } impl CacheHttp for CacheAndHttp { fn http(&self) -> &Http { &self.1 } fn cache(&self) -> Option<&Arc> { Some(&self.0) } } /// A list of assigned channels for an announcer. pub(crate) type AnnouncerChannels = DB>>; /// The Announcer trait. /// /// Every announcer needs to implement a method to look for updates. /// This method is called "updates", which takes: /// - A CacheHttp implementation, for interaction with Discord itself. /// - An AppData, which can be used for interacting with internal databases. /// - A function "channels", which takes an UserId and returns the list of ChannelIds, which any update related to that user should be /// sent to. #[async_trait] pub trait Announcer: Send { /// Look for updates and send them to respective channels. /// /// Errors returned from this function gets ignored and logged down. async fn updates( &mut self, c: CacheAndHttp, d: AppData, channels: MemberToChannels, ) -> Result<()>; } /// A simple struct that allows looking up the relevant channels to an user. pub struct MemberToChannels(Vec<(GuildId, ChannelId)>, AppData); impl MemberToChannels { /// Gets the channel list of an user related to that channel. pub async fn channels_of( &self, http: impl CacheHttp + Clone, u: impl Into, ) -> Vec { let u: UserId = u.into(); let member_cache = self.1.read().await.get::().unwrap().clone(); self.0 .clone() .into_iter() .map(|(guild, channel)| { member_cache .query(http.clone(), u, guild) .map(move |t| t.map(|_| channel)) }) .collect::>() .filter_map(ready) .collect() .await } } /// The announcer handler. /// /// This struct manages the list of all Announcers, firing them in a certain interval. pub struct AnnouncerHandler { announcers: HashMap<&'static str, RwLock>>, } /// Announcer-managing related. impl Default for AnnouncerHandler { fn default() -> Self { Self::new() } } impl AnnouncerHandler { /// Create a new instance of the handler. pub fn new() -> Self { Self { announcers: HashMap::new(), } } /// Insert a new announcer into the handler. /// /// The handler must take an unique key. If a duplicate is found, this method panics. pub fn add( &mut self, key: &'static str, announcer: impl Announcer + Sync + 'static, ) -> &mut Self { if self .announcers .insert(key, RwLock::new(Box::new(announcer))) .is_some() { panic!( "Announcer keys must be unique: another announcer with key `{}` was found", key ) } else { self } } pub fn run(self, client: &Client) -> AnnouncerRunner { AnnouncerRunner { cache_http: CacheAndHttp::from_client(client), data: client.data.clone(), announcers: self.announcers, } } } pub struct AnnouncerRunner { cache_http: CacheAndHttp, data: AppData, announcers: HashMap<&'static str, RwLock>>, } // Querying for the AnnouncerRunner in the internal data returns a vec of keys. impl TypeMapKey for AnnouncerRunner { type Value = Vec<&'static str>; } /// Execution-related. impl AnnouncerRunner { /// Collect the list of guilds and their respective channels, by the key of the announcer. async fn get_guilds(data: &AppData, key: &'static str) -> Result> { let data = AnnouncerChannels::open(&*data.read().await) .borrow()? .get(key) .map(|m| m.iter().map(|(a, b)| (*a, *b)).collect()) .unwrap_or_else(Vec::new); Ok(data) } /// Run the announcing sequence on a certain announcer. async fn announce( data: AppData, cache_http: CacheAndHttp, key: &'static str, announcer: &'_ RwLock>, ) -> Result<()> { let channels = MemberToChannels(Self::get_guilds(&data, key).await?, data.clone()); announcer .write() .await .updates(cache_http, data, channels) .await } /// Start the AnnouncerHandler, looping forever. /// /// It will run all the announcers every *cooldown* seconds. pub async fn scan(self, cooldown: std::time::Duration) { // First we store all the keys inside the database. let keys = self.announcers.keys().cloned().collect::>(); self.data.write().await.insert::(keys.clone()); join_all(self.announcers.iter().map(|(key, announcer)| { let data = self.data.clone(); let cache = self.cache_http.clone(); async move { loop { eprintln!(" - scanning key `{}`", key); match Self::announce(data.clone(), cache.clone(), key, announcer).await { Err(e) => { eprintln!(" - key `{}`: {:?}", *key, e) } Ok(()) => { eprintln!(" - key `{}`: complete", *key) } }; tokio::time::sleep(cooldown).await; } } })) .await; } } /// Gets the announcer of the given guild. pub async fn announcer_of( ctx: &Context, key: &'static str, guild: GuildId, ) -> Result> { Ok(AnnouncerChannels::open(&*ctx.data.read().await) .borrow()? .get(key) .and_then(|channels| channels.get(&guild).cloned())) } #[command("list")] #[description = "List the registered announcers of this server"] #[num_args(0)] #[only_in(guilds)] pub async fn list_announcers(ctx: &Context, m: &Message, _: Args) -> CommandResult { let guild_id = m.guild_id.unwrap(); let data = &*ctx.data.read().await; let announcers = AnnouncerChannels::open(data); let channels = data.get::().unwrap(); let channels = channels .iter() .filter_map(|&key| { announcers.borrow().ok().and_then(|announcers| { announcers .get(key) .and_then(|channels| channels.get(&guild_id)) .map(|&ch| (key, ch)) }) }) .map(|(key, ch)| format!(" - `{}`: activated on channel {}", key, ch.mention())) .collect::>(); m.reply( &ctx, format!( "Activated announcers on this server:\n{}", channels.join("\n") ), ) .await?; Ok(()) } #[command("register")] #[description = "Register the current channel with an announcer"] #[usage = "[announcer key]"] #[required_permissions(MANAGE_CHANNELS)] #[only_in(guilds)] #[num_args(1)] pub async fn register_announcer(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let data = ctx.data.read().await; let key = args.single::()?; let keys = data.get::().unwrap(); if !keys.contains(&&key[..]) { m.reply( &ctx, format!( "Key not found. Available announcer keys are: `{}`", keys.join(", ") ), ) .await?; return Ok(()); } let guild = m .guild_id .expect("Guild-only command") .to_partial_guild(&ctx) .await?; let channel = m.channel_id.to_channel(&ctx).await?; AnnouncerChannels::open(&data) .borrow_mut()? .entry(key.clone()) .or_default() .insert(guild.id, m.channel_id); m.reply( &ctx, MessageBuilder::new() .push("Announcer ") .push_mono_safe(key) .push(" has been activated for server ") .push_bold_safe(&guild.name) .push(" on channel ") .push_bold_safe(channel.mention().to_string()) .build(), ) .await?; Ok(()) } #[command("remove")] #[description = "Remove an announcer from the server"] #[usage = "[announcer key]"] #[required_permissions(MANAGE_CHANNELS)] #[only_in(guilds)] #[num_args(1)] pub async fn remove_announcer(ctx: &Context, m: &Message, mut args: Args) -> CommandResult { let data = ctx.data.read().await; let key = args.single::()?; let keys = data.get::().unwrap(); if !keys.contains(&key.as_str()) { m.reply( &ctx, format!( "Key not found. Available announcer keys are: `{}`", keys.join(", ") ), ) .await?; return Ok(()); } let guild = m .guild_id .expect("Guild-only command") .to_partial_guild(&ctx) .await?; AnnouncerChannels::open(&data) .borrow_mut()? .entry(key.clone()) .and_modify(|m| { m.remove(&guild.id); }); m.reply( &ctx, MessageBuilder::new() .push("Announcer ") .push_mono_safe(key) .push(" has been de-activated for server ") .push_bold_safe(&guild.name) .build(), ) .await?; Ok(()) } #[group("announcer")] #[prefix("announcer")] #[only_in(guilds)] #[required_permissions(MANAGE_CHANNELS)] #[description = "Manage the announcers in the server."] #[commands(remove_announcer, register_announcer, list_announcers)] pub struct AnnouncerCommands;