use std::{future::Future, path::Path}; use ::serenity::all::{ChannelId, GuildId, UserId}; use eyre::{Context as _, Error, OptionExt, Result}; use futures::{FutureExt, StreamExt, TryStreamExt}; use poise::serenity_prelude as serenity; use rand::seq::SliceRandom; use rusqlite::OptionalExtension; #[derive(Copy, Clone)] struct GuildData { guild: serenity::GuildId, town_square: serenity::ChannelId, cottages: serenity::ChannelId, st: serenity::RoleId, currently_playing: serenity::RoleId, } struct Data(tokio::sync::Mutex); impl Data { fn open(path: impl AsRef) -> Result { let conn = rusqlite::Connection::open(path)?; conn.execute_batch("CREATE TABLE IF NOT EXISTS guilds(guild TEXT PRIMARY KEY NOT NULL, town_square TEXT NOT NULL, cottages TEXT NOT NULL, st TEXT, currently_playing TEXT);")?; Ok(Self(tokio::sync::Mutex::new(conn))) } async fn get(&self, guild: serenity::GuildId) -> Result { self.0 .lock() .await .prepare_cached("SELECT guild, town_square, cottages, st, currently_playing FROM guilds WHERE guild = ?;")? .query_row(rusqlite::params![guild.to_string()], |row| { Ok(GuildData { guild: row.get_unwrap::<_, String>("guild").parse().unwrap(), town_square: row.get_unwrap::<_, String>("town_square").parse().unwrap(), cottages: row.get_unwrap::<_, String>("cottages").parse().unwrap(), st: row.get_unwrap::<_, String>("st").parse().unwrap(), currently_playing: row.get_unwrap::<_, String>("currently_playing").parse().unwrap(), }) }) .optional()? .ok_or_eyre("This bot has not been configured in this server!") } async fn insert( &self, GuildData { guild, town_square, cottages, st, currently_playing, }: GuildData, ) -> Result<()> { self.0 .lock() .await .prepare_cached("INSERT INTO guilds (guild, town_square, cottages, st, currently_playing) VALUES (?, ?, ?, ?, ?) ON CONFLICT(guild) DO UPDATE SET town_square = excluded.town_square, cottages = excluded.cottages, st = excluded.st, currently_playing = excluded.currently_playing; ")? .execute(rusqlite::params![ guild.to_string(), town_square.to_string(), cottages.to_string(), st.to_string(), currently_playing.to_string(), ])?; Ok(()) } } type Context<'a> = poise::Context<'a, Data, Error>; async fn channel_children( ctx: &Context<'_>, channel: serenity::ChannelId, ) -> Result> { let channels = ctx .guild_id() .ok_or_eyre("This bot only works in servers")? .channels(&ctx) .await?; Ok(channels .into_values() .filter(|chan| chan.parent_id == Some(channel)) .collect()) } fn move_users<'a>( ctx: Context<'a>, guild: GuildId, users: impl Iterator + Send + 'a, ) -> impl Future> + Send + 'a { futures::stream::iter(users.map(move |(channel, user)| async move { guild.move_member(ctx, user, channel).await?; Ok(()) })) .buffer_unordered(10) .try_collect::<()>() } async fn spread( ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.defer_ephemeral().await?; let from = from .to_channel(ctx) .await? .guild() .ok_or_eyre("This bot only works in servers")?; let mut to = channel_children(&ctx, to).await?; let mut users = from.members(ctx)?; users.shuffle(&mut rand::thread_rng()); to.retain(|x| x.members(ctx).is_ok_and(|x| x.is_empty())); to.shuffle(&mut rand::thread_rng()); move_users( ctx, guild, to.into_iter() .zip(&users) .map(|(channel, user)| (channel.id, user.user.id)), ) .await?; ctx.reply(format!("Moved {} users", users.len())).await?; Ok(()) } async fn collect( ctx: Context<'_>, from: serenity::ChannelId, to: serenity::ChannelId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.defer_ephemeral().await?; let users = channel_children(&ctx, from) .await? .into_iter() .flat_map(|x| x.members(ctx)) .flatten() .collect::>(); move_users(ctx, guild, users.iter().map(|user| (to, user.user.id))).await?; futures::future::try_join_all(users.iter().map(|user| guild.move_member(&ctx, user, to))) .await?; ctx.reply(format!("Moved {} users", users.len())).await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] async fn dusk(ctx: Context<'_>) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { town_square, cottages, .. } = ctx.data().get(guild).await?; spread(ctx, town_square, cottages).await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] async fn dawn(ctx: Context<'_>) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { town_square, cottages, .. } = ctx.data().get(guild).await?; collect(ctx, cottages, town_square).await?; Ok(()) } #[poise::command(slash_command, ephemeral)] async fn join( ctx: Context<'_>, to: serenity::UserId, user: Option, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { cottages, .. } = ctx.data().get(guild).await?; let channel = channel_children(&ctx, cottages) .await? .into_iter() .find(|chan| { chan.members(ctx) .is_ok_and(|members| members.iter().any(|member| member.user.id == to)) }); match channel { Some(channel) => { guild .move_member(ctx, user.unwrap_or(ctx.author().id), channel) .await?; ctx.reply("Joined player").await?; } None => { ctx.reply("User is not in cottages").await?; } } Ok(()) } #[poise::command(slash_command, ephemeral)] async fn st(ctx: Context<'_>, mut spectators: Vec) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { st, .. } = ctx.data().get(guild).await?; spectators.push(ctx.author().id); futures::future::try_join_all(guild.members(&ctx, None, None).await?.into_iter().map( |member| { let spectators = &spectators; async move { match ( spectators.contains(&member.user.id), member.roles.contains(&st), ) { (true, true) | (false, false) => Ok(()), (true, false) => member.add_role(&ctx, st).await, (false, true) => member.remove_role(&ctx, st).await, } } }, )) .await?; ctx.reply(format!("{} members given ST", spectators.len())) .await?; Ok(()) } #[poise::command(slash_command, ephemeral)] async fn currently_playing(ctx: Context<'_>, players: Vec) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; let GuildData { currently_playing, .. } = ctx.data().get(guild).await?; for player in &players { guild .member(&ctx, player) .await? .add_role(&ctx, currently_playing) .await?; } ctx.reply(format!("Given {} members currently playing", players.len())) .await?; Ok(()) } #[poise::command(slash_command, prefix_command, ephemeral)] async fn configure( ctx: Context<'_>, #[channel_types("Voice")] town_square: serenity::ChannelId, #[channel_types("Category")] cottages: serenity::ChannelId, st: serenity::RoleId, currently_playing: serenity::RoleId, ) -> Result<()> { let guild = ctx .guild_id() .ok_or_eyre("This bot only works in servers")?; ctx.data() .insert(GuildData { guild, town_square, cottages, st, currently_playing, }) .await?; ctx.reply("Configured successfully!").await?; Ok(()) } async fn register_commands( ctx: impl AsRef, options: &poise::FrameworkOptions, guild: serenity::GuildId, ) -> Result<()> { poise::builtins::register_in_guild(ctx, &options.commands, guild).await?; Ok(()) } async fn on_event( ctx: &serenity::Context, event: &serenity::FullEvent, framework: poise::FrameworkContext<'_, Data, Error>, _: &Data, ) -> Result<()> { if let serenity::FullEvent::GuildCreate { guild, is_new: _ } = event { register_commands(ctx, framework.options, guild.id).await?; }; Ok(()) } #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { color_eyre::install()?; let token = std::fs::read_to_string( std::env::var_os("DISCORD_TOKEN_FILE").ok_or_eyre("A Discord token is required")?, ) .wrap_err("A Discord token is required")?; let dbfile = std::env::var_os("DB_FILE").ok_or_eyre("A database file is required")?; let intents = serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::GUILD_MEMBERS; let framework = poise::Framework::builder() .options(poise::FrameworkOptions { commands: vec![ dusk(), dawn(), st(), currently_playing(), join(), configure(), ], event_handler: |ctx, ev, ctxf, data| on_event(ctx, ev, ctxf, data).boxed(), ..Default::default() }) .setup(|ctx, ready, framework| { Box::pin(async move { for guild in &ready.guilds { register_commands(ctx, framework.options(), guild.id).await?; } Data::open(dbfile) }) }) .build(); serenity::ClientBuilder::new(token, intents) .framework(framework) .await? .start() .await?; Ok(()) }