nat1/src/main.rs

use std::fmt::Write as FmtWrite;
mod dice;
mod limits;
mod parsing;

use crate::dice::{DiscordMd, Witness};
use anyhow::{Context as AnyCtx, bail};
use dotenv::dotenv;
use limits::*;
use rand::{Rng, SeedableRng};
use serenity::all::{
    Command, CommandInteraction, CommandOptionType, CreateCommandOption, CreateComponent,
    FullEvent, Interaction, InteractionContext, MessageFlags, ResolvedOption, ResolvedValue,
};
use serenity::builder::{
    CreateCommand, CreateContainer, CreateContainerComponent, CreateInteractionResponse,
    CreateInteractionResponseMessage, CreateTextDisplay,
};
use serenity::{async_trait, prelude::*};
use std::sync::Arc;
use tracing::{debug, info, warn};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    dotenv().ok();
    tracing_subscriber::fmt::init();

    // Configure the client with your Discord bot token in the environment.
    let token =
        Token::from_env("DISCORD_TOKEN").expect("Expected a valid token in the environment");

    // Build our client.
    let mut client = Client::builder(token, GatewayIntents::empty())
        .event_handler(Arc::new(Handler))
        .await
        .expect("Error creating client");

    client.start().await?;
    Ok(())
}

struct Handler;

#[async_trait]
impl EventHandler for Handler {
    async fn dispatch(&self, ctx: &Context, event: &FullEvent) {
        // clippy can't decide between if it wants it collapsed, or if it wants you to use if let
        // because it's a single pattern.
        #[expect(clippy::collapsible_match)]
        match event {
            FullEvent::InteractionCreate { interaction, .. } => {
                if let Interaction::Command(command) = interaction {
                    debug!("Received command interaction: {command:#?}");

                    let content = match command.data.name.as_str() {
                        "roll" => {
                            roll(ctx, command)
                                .await
                                .unwrap_or_else(|e| warn!("error! {}", e));
                            None
                        }
                        _ => Some("not implemented :(".to_string()),
                    };

                    if let Some(content) = content {
                        let data = CreateInteractionResponseMessage::new().content(content);
                        let builder = CreateInteractionResponse::Message(data);
                        if let Err(why) = command.create_response(&ctx.http, builder).await {
                            warn!("Cannot respond to slash command: {why}");
                        }
                    }
                }
            }
            FullEvent::Ready { data_about_bot, .. } => {
                let global_command = Command::create_global_command(
                    &ctx.http,
                    CreateCommand::new("roll")
                        .description("Roll some dice!")
                        .add_option(CreateCommandOption::new(
                            CommandOptionType::String,
                            "formula",
                            "The dice formula, something like 2d20kh1+3",
                        ).required(true))
                        .add_option(CreateCommandOption::new(
                            CommandOptionType::Integer,
                            "repeat",
                            "Repeat the rolled formula this many times. Useful if you're controlling too many units!",
                        ))
                        .add_option(CreateCommandOption::new(
                            CommandOptionType::String,
                            "description",
                            "What is that roll for? Add some context!",
                        ))
                        .add_option(CreateCommandOption::new(
                            CommandOptionType::Boolean,
                            "private",
                            "Keep the roll private and hidden, only for yourself.",
                        ))
                        .add_option(CreateCommandOption::new(
                            CommandOptionType::String,
                            "seed",
                            "Roll with a fixed seed. This bot uses xoshiro256++.",
                        ))
                        .add_context(InteractionContext::Guild)
                        .add_context(InteractionContext::BotDm)
                        .add_context(InteractionContext::PrivateChannel),
                )
                    .await;

                debug!("I created the following global slash command: {global_command:#?}");

                info!("{} is connected!", data_about_bot.user.name);
            }
            _ => {}
        }
    }
}

pub async fn roll(ctx: &Context, command: &CommandInteraction) -> anyhow::Result<()> {
    pub async fn inner(ctx: &Context, command: &CommandInteraction) -> anyhow::Result<()> {
        let options = command.data.options();
        let Some(ResolvedOption {
            value: ResolvedValue::String(formula),
            ..
        }) = options.iter().find(|option| option.name == "formula")
        else {
            bail!("Formula was not provided.")
        };
        let repeat = options
            .iter()
            .find(|option| option.name == "repeat")
            .and_then(|option| {
                if let ResolvedValue::Integer(repeat) = option.value {
                    Some(repeat)
                } else {
                    None
                }
            })
            .unwrap_or(1)
            .clamp(1, MAX_ROLL_REPEATS as i64) as usize;
        let private = options
            .iter()
            .find(|option| option.name == "private")
            .and_then(|option| {
                if let ResolvedValue::Boolean(private) = option.value {
                    Some(private)
                } else {
                    None
                }
            })
            .unwrap_or(false);
        let description = options
            .iter()
            .find(|option| option.name == "description")
            .and_then(|option| {
                if let ResolvedValue::String(description) = option.value {
                    Some(description)
                } else {
                    None
                }
            });
        let fixed_seed = options
            .iter()
            .find(|option| option.name == "seed")
            .and_then(|option| {
                if let ResolvedValue::String(fixed_seed) = option.value {
                    Some(fixed_seed)
                } else {
                    None
                }
            })
            .map(|s| s.parse::<u64>().context("Seed could not be parsed. Make sure it is a 64 bit integer between 0 and 2^64-1."))
            .transpose()?;

        let expression = parsing::parse(formula).context("Dice formula could not be parsed.")?;

        let seed = fixed_seed.unwrap_or_else(|| rand::rng().next_u64());
        let mut rng = rand_xoshiro::Xoshiro256PlusPlus::seed_from_u64(seed);

        let average = expression.average(&mut rng)?;
        let ev = (average * 10.0).round() / 10.0;

        let mut pre_paragraphs = Vec::new();
        if let Some(description) = description {
            pre_paragraphs.insert(0, format!(">>> {description}\n"));
        }

        let mut post_paragraphs = Vec::new();
        if fixed_seed.is_some() {
            post_paragraphs.push("***⚠️  This is not a fair roll. The result was generated based on a fixed seed and is thus fully reproducible.***".to_string());
        }
        post_paragraphs.push(format!("-# Expected value: {ev}  |  Seed: {seed}"));

        let surrounding_par_count = pre_paragraphs.len() + post_paragraphs.len();

        let roll_paragraphs = if repeat <= 40 - surrounding_par_count {
            let discord_md = Arc::new(std::sync::RwLock::new(DiscordMd {
                buffers: Vec::with_capacity(repeat),
                dice_written: 0,
            }));
            for _ in 0..repeat {
                let mut witness = DiscordMd::create_witness(discord_md.clone());
                let result = expression.evaluate(&mut rng, &mut witness, i64::MAX)?;
                witness.witness_total_result(result)?;
                witness.end();
            }
            let discord_md = Arc::into_inner(discord_md).unwrap().into_inner()?;

            discord_md.buffers
        } else {
            let mut par = String::with_capacity(8 * repeat);
            for i in 0..repeat {
                if i > 0 {
                    write!(par, ", ")?;
                }
                write!(
                    par,
                    "**{}**",
                    expression.evaluate(&mut rng, &mut std::io::sink(), i64::MAX)?
                )?;
            }
            vec![par]
        };

        let paragraphs = {
            pre_paragraphs.extend(roll_paragraphs);
            pre_paragraphs.extend(post_paragraphs);
            pre_paragraphs
        };

        if paragraphs.iter().map(|s| s.chars().count()).sum::<usize>() > 4000
            || paragraphs.len() > 40
        {
            bail!("Resulting message was too long, please roll fewer dice!");
        };

        let text_components = paragraphs
            .into_iter()
            .map(|text| CreateContainerComponent::TextDisplay(CreateTextDisplay::new(text)))
            .collect::<Vec<_>>();

        let response = CreateInteractionResponse::Message(
            CreateInteractionResponseMessage::new()
                .components(vec![CreateComponent::Container(CreateContainer::new(
                    text_components,
                ))])
                // flags needs to be called before ephemeral for correct ordering
                .flags(MessageFlags::IS_COMPONENTS_V2 | MessageFlags::SUPPRESS_NOTIFICATIONS)
                .ephemeral(private),
        );

        debug!("Returning response: {response:#?}");

        command.create_response(&ctx.http, response).await?;
        Ok(())
    }

    match inner(ctx, command).await {
        Ok(_) => {}
        Err(err) => {
            warn!("Error encountered: {err:?}");
            command
                .create_response(
                    &ctx.http,
                    CreateInteractionResponse::Message(
                        CreateInteractionResponseMessage::new()
                            .components(vec![CreateComponent::Container(CreateContainer::new(
                                vec![CreateContainerComponent::TextDisplay(
                                    CreateTextDisplay::new(err.to_string()),
                                )],
                            ))])
                            // flags needs to be called before ephemeral for correct ordering
                            .flags(
                                MessageFlags::IS_COMPONENTS_V2 | MessageFlags::SUPPRESS_NOTIFICATIONS,
                            )
                            .ephemeral(true),
                    ),
                )
                .await?
        }
    };
    Ok(())
}