reminder-bot/src/main.rs
2024-02-09 17:03:04 +00:00

271 lines
8.2 KiB
Rust

#![feature(int_roundings)]
#[macro_use]
extern crate lazy_static;
mod commands;
mod component_models;
mod consts;
mod event_handlers;
mod hooks;
mod interval_parser;
mod models;
mod time_parser;
mod utils;
use std::{
collections::HashMap,
env,
error::Error as StdError,
fmt::{Debug, Display, Formatter},
path::Path,
};
use chrono_tz::Tz;
use log::{error, warn};
use poise::serenity_prelude::{
model::{
gateway::GatewayIntents,
id::{GuildId, UserId},
},
ClientBuilder,
};
use sqlx::{MySql, Pool};
use tokio::sync::{broadcast, broadcast::Sender, RwLock};
use crate::{
commands::{command_macro, info_cmds, moderation_cmds, reminder_cmds, todo_cmds},
consts::THEME_COLOR,
event_handlers::listener,
hooks::all_checks,
models::command_macro::CommandMacro,
};
type Database = MySql;
type Error = Box<dyn std::error::Error + Send + Sync>;
type Context<'a> = poise::Context<'a, Data, Error>;
type ApplicationContext<'a> = poise::ApplicationContext<'a, Data, Error>;
pub struct Data {
database: Pool<Database>,
recording_macros: RwLock<HashMap<(GuildId, UserId), CommandMacro>>,
popular_timezones: Vec<Tz>,
_broadcast: Sender<()>,
}
impl Debug for Data {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "Data {{ .. }}")
}
}
struct Ended;
impl Debug for Ended {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str("Process ended.")
}
}
impl Display for Ended {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
f.write_str("Process ended.")
}
}
impl StdError for Ended {}
#[tokio::main(flavor = "multi_thread")]
async fn main() -> Result<(), Box<dyn StdError + Send + Sync>> {
let (tx, mut rx) = broadcast::channel(16);
tokio::select! {
output = _main(tx) => output,
_ = rx.recv() => Err(Box::new(Ended) as Box<dyn StdError + Send + Sync>)
}
}
async fn _main(tx: Sender<()>) -> Result<(), Box<dyn StdError + Send + Sync>> {
env_logger::init();
if Path::new("/etc/reminder-rs/config.env").exists() {
dotenv::from_path("/etc/reminder-rs/config.env")?;
} else {
let _ = dotenv::dotenv();
}
let discord_token = env::var("DISCORD_TOKEN").expect("Missing DISCORD_TOKEN from environment");
let options = poise::FrameworkOptions {
commands: vec![
info_cmds::help(),
info_cmds::info(),
info_cmds::donate(),
info_cmds::clock(),
info_cmds::clock_context_menu(),
info_cmds::dashboard(),
moderation_cmds::timezone(),
poise::Command {
subcommands: vec![
moderation_cmds::set_allowed_dm(),
moderation_cmds::unset_allowed_dm(),
],
..moderation_cmds::allowed_dm()
},
poise::Command {
subcommands: vec![poise::Command {
subcommands: vec![
moderation_cmds::set_ephemeral_confirmations(),
moderation_cmds::unset_ephemeral_confirmations(),
],
..moderation_cmds::ephemeral_confirmations()
}],
..moderation_cmds::settings()
},
moderation_cmds::webhook(),
poise::Command {
subcommands: vec![
command_macro::delete::delete_macro(),
command_macro::record::finish_macro(),
command_macro::list::list_macro(),
command_macro::record::record_macro(),
command_macro::run::run_macro(),
],
..command_macro::macro_base()
},
reminder_cmds::pause(),
reminder_cmds::offset(),
reminder_cmds::nudge(),
reminder_cmds::look(),
reminder_cmds::delete(),
poise::Command {
subcommands: vec![
reminder_cmds::list_timer(),
reminder_cmds::start_timer(),
reminder_cmds::delete_timer(),
],
..reminder_cmds::timer_base()
},
reminder_cmds::multiline(),
reminder_cmds::remind(),
poise::Command {
subcommands: vec![
poise::Command {
subcommands: vec![
todo_cmds::todo_guild_add(),
todo_cmds::todo_guild_view(),
],
..todo_cmds::todo_guild_base()
},
poise::Command {
subcommands: vec![
todo_cmds::todo_channel_add(),
todo_cmds::todo_channel_view(),
],
..todo_cmds::todo_channel_base()
},
poise::Command {
subcommands: vec![todo_cmds::todo_user_add(), todo_cmds::todo_user_view()],
..todo_cmds::todo_user_base()
},
],
..todo_cmds::todo_base()
},
],
allowed_mentions: None,
command_check: Some(|ctx| Box::pin(all_checks(ctx))),
event_handler: |ctx, event, _framework, data| Box::pin(listener(ctx, event, data)),
on_error: |error| {
Box::pin(async move {
match error {
poise::FrameworkError::CommandCheckFailed { .. } => {
// suppress error
}
error => {
if let Err(e) = poise::builtins::on_error(error).await {
log::error!("Error while handling error: {}", e);
}
}
}
})
},
..Default::default()
};
let database =
Pool::connect(&env::var("DATABASE_URL").expect("No database URL provided")).await.unwrap();
sqlx::migrate!().run(&database).await?;
let popular_timezones = sqlx::query!(
"SELECT IFNULL(timezone, 'UTC') AS timezone
FROM users
WHERE timezone IS NOT NULL
GROUP BY timezone
ORDER BY COUNT(timezone) DESC
LIMIT 21"
)
.fetch_all(&database)
.await
.unwrap()
.iter()
.map(|t| t.timezone.parse::<Tz>().unwrap())
.collect::<Vec<Tz>>();
let framework = poise::Framework::builder()
.setup(move |ctx, _bot, framework| {
Box::pin(async move {
poise::builtins::register_globally(ctx, &framework.options().commands).await?;
let kill_tx = tx.clone();
let kill_recv = tx.subscribe();
let ctx1 = ctx.clone();
let ctx2 = ctx.clone();
let pool1 = database.clone();
let pool2 = database.clone();
let run_settings = env::var("DONTRUN").unwrap_or_else(|_| "".to_string());
if !run_settings.contains("postman") {
tokio::spawn(async move {
match postman::initialize(kill_recv, ctx1, &pool1).await {
Ok(_) => {}
Err(e) => {
error!("postman exiting: {}", e);
}
};
});
} else {
warn!("Not running postman");
}
if !run_settings.contains("web") {
tokio::spawn(async move {
reminder_web::initialize(kill_tx, ctx2, pool2).await.unwrap();
});
} else {
warn!("Not running web");
}
Ok(Data {
database,
popular_timezones,
recording_macros: Default::default(),
_broadcast: tx,
})
})
})
.options(options)
.build();
let mut client =
ClientBuilder::new(&discord_token, GatewayIntents::GUILDS).framework(framework).await?;
client.start_autosharded().await?;
Ok(())
}