botc-mover/src/main.rs

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<rusqlite::Connection>);

impl Data {
    fn open(path: impl AsRef<Path>) -> Result<Self> {
        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<GuildData> {
        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<Vec<serenity::GuildChannel>> {
    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<Item = (ChannelId, UserId)> + Send + 'a,
) -> impl Future<Output = Result<()>> + 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::<Vec<_>>();

    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<serenity::UserId>,
) -> 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<serenity::UserId>) -> 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<serenity::UserId>) -> 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<serenity::Http>,
    options: &poise::FrameworkOptions<Data, Error>,
    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(())
}