path: root/src/commands/reactionroles.rs

use std::{str::FromStr, time::Duration};

use crate::{models::ReactionRole, Context, Error};
use ::serenity::{
    model::{guild::Role, id::ChannelId},
};
use poise::serenity_prelude::{self as serenity, ArgumentConvert, Emoji, ReactionType, RoleId};

#[cfg(debug_assertions)]
async fn allowed_to_create_roles(ctx: Context<'_>) -> Result<bool, Error> {
    if ctx.author().id.0 == 118455061222260736u64 {
        Ok(true)
    } else {
        Ok(false)
    }
}

#[cfg(not(debug_assertions))]
async fn allowed_to_create_roles(ctx: Context<'_>) -> Result<bool, Error> {
    if let Some(guild) = ctx.guild() {
        if guild.owner_id == ctx.author().id {
            Ok(true)
        } else {
            let member = guild.member(ctx.discord(), ctx.author().id).await?;
            let member_permissions = member.permissions(ctx.discord())?;

            Ok(member_permissions.manage_roles())
        }
    } else {
        Ok(false)
    }
}

/// Manages reaction role menus
///
/// Subcommands:
///     - init
///     - add
///     - remove
#[poise::command(prefix_command, ephemeral, check = "allowed_to_create_roles")]
pub async fn rroles(ctx: Context<'_>) -> Result<(), Error> {
    ctx.say("Maybe you meant to request help for this? Try `/help rroles`")
        .await?;
    Ok(())
}

/// Initializes a reaction role menu in the given channel
///
/// Usage:
///     rroles init <#channel>
/// Example:
///     rroles init #get-roles
#[poise::command(prefix_command, ephemeral, check = "allowed_to_create_roles")]
pub async fn init(
    ctx: Context<'_>,
    #[description = "The channel to create a new role menu in"] channel: serenity::ChannelId,
) -> Result<(), Error> {
    let mut rolemenu_msg = channel
        .send_message(ctx.discord(), |m| {
            m.embed(|e| {
                e.title("Reaction Role Menu");
                e.description("I haven't been initialized yet! Hold on just a second");
                e
            });
            m
        })
        .await?;

    let mut menu = ctx.send(|m| {
        m.embed(|e| {
            e.title("Reaction Role Setup");
            e.description("Welcome to the setup menu! I'm going to help guide you through setting up your reaction roles.\n\nFirst, what should the title of your menu?");
            e
        });
        m
    }).await?.unwrap().message().await?;

    if let Some(title) = ctx
        .author()
        .clone()
        .await_reply(ctx.discord())
        .timeout(Duration::from_secs(10))
        .await
    {
        rolemenu_msg
            .edit(ctx.discord(), |m| {
                m.embed(|e| {
                    e.title(title.content.clone());
                    e
                });
                m
            })
            .await?;

        menu.edit(ctx.discord(), |m| {
            m.embed(|e| {
                e.title("Reaction Role Setup");
                e.description(format!("Great! I've set the title of your menu to `{}`.\n\nNext, let's add some roles! Send the first emoji you want to add.", title.content.clone()));
                e
            });
            m
        }).await?;
    } else {
        ctx.say("No response within 10 seconds").await?;
        return Ok(());
    }
    {
        let pool = ctx.data().pg.lock().unwrap().clone();
        loop {
            if let Some(emoji) = ctx
                .author()
                .clone()
                .await_reply(ctx.discord())
                .timeout(Duration::from_secs(30))
                .await
            {
                let content = emoji.content.clone();
                if content == "done" {
                    menu.edit(ctx.discord(), |m| {
                        m.embed(|e| {
                        e.title("Reaction Role Setup");
                        e.description("Nice work! You're all set up! Use `~rroles add` and `~rroles del` to manage the roles in this menu!");
                        e
                    });
                    m
                }).await?;
                    break;
                }
                let reaction = ReactionType::from_str(&emoji.content)?;

                menu.edit(ctx.discord(), |m| {
                    m.embed(|e| {
                        e.title("Reaction Role Setup");
                        e.description(format!("Sounds good! Let's give {} a role, okay? Reply to this message with the name of the role you'd like to assign to this emoji", reaction.clone()));
                        e
                    });
                    m
                }).await?;

                if let Some(role) = ctx
                    .author()
                    .clone()
                    .await_reply(ctx.discord())
                    .timeout(Duration::from_secs(30))
                    .await
                {
                    if let Some(role) = ctx.guild().unwrap().role_by_name(&role.content) {
                        sqlx::query!("INSERT INTO reaction_roles (channel_id, message_id, guild_id, reaction, role_id) VALUES ($1, $2, $3, $4, $5)", rolemenu_msg.channel_id.0.to_string(), rolemenu_msg.id.0.to_string(), ctx.guild_id().unwrap().0.to_string(), reaction.to_string(), role.id.0.to_string()).execute(&pool).await?;
                        rolemenu_msg.react(ctx.discord(), reaction.clone()).await?;

                        menu.edit(ctx.discord(), |m| {
            m.embed(|e| {
                e.title("Reaction Role Setup");
                e.description(format!("Great! I've added that to the menu.\n\nLet's keep adding roles! Send the next emoji you want to add, or 'done' to finish setup."));
                e
            });
            m
        }).await?;
                    } else {
                        menu.edit(ctx.discord(), |m| {
                            m.embed(|e| {
                                e.title("Reaction Role Setup");
                                e.description("Whoops! I couldn't find that role! Let's try again! Send the emoji you want to assign to a role");
                                e
                            });
                            m
                        }).await?;
                        continue;
                    }
                }

                let reactions = sqlx::query_as!(
                    ReactionRole,
                    "SELECT * FROM reaction_roles WHERE message_id=$1",
                    rolemenu_msg.id.0.to_string()
                )
                .fetch_all(&pool)
                .await?;

                let rolelist_formatted = gen_reaction_list(reactions);

                let title = rolemenu_msg.clone().embeds[0]
                    .title
                    .clone()
                    .unwrap_or("Reaction Role Menu".to_string());
                rolemenu_msg
                    .edit(ctx.discord(), |m| {
                        m.embed(|e| {
                            e.title(title);
                            e.description(rolelist_formatted);
                            e
                        });
                        m
                    })
                    .await?;
            } else {
                ctx.say("No response within 30 seconds").await?;
                return Ok(());
            }
        }
    }

    Ok(())
}

/// Adds a reaction role to the message
///
/// Usage:
///     ~rroles add <Message ID> <emoji> <role>
#[poise::command(prefix_command, check = "allowed_to_create_roles")]
pub async fn add(
    ctx: Context<'_>,
    #[description = "The Message ID"] message_id: u64,
    #[description = "The emoji to assign to the role"] emoji: ReactionType,
    #[description = "The role to assign to the emoji"] role_name: String,
) -> Result<(), Error> {
    {
        let pool = ctx.data().pg.lock().unwrap().clone();
        // Make sure the emoji doesn't already exist
        if let Ok(_) = sqlx::query!("SELECT * FROM reaction_roles WHERE message_id=$1 AND reaction=$2", message_id.to_string(), emoji.to_string()).fetch_one(&pool).await {
            ctx.say("Whoops! That emoji already has something assigned to it! Try either removing it or picking a different emoji").await?;
            return Ok(());
        }
        let role_menu = sqlx::query_as!(
            ReactionRole,
            "SELECT * FROM reaction_roles WHERE message_id=$1",
            message_id.to_string()
        )
        .fetch_one(&pool)
        .await?;
        let guild = ctx.guild().unwrap();
        let role = guild.role_by_name(&role_name).clone();
        if let Some(r) = role.clone() {
            let r = r.clone();
            let channel_id = ChannelId(role_menu.channel_id.parse::<u64>()?);
            sqlx::query!("INSERT INTO reaction_roles (channel_id, message_id, guild_id, reaction, role_id) VALUES ($1, $2, $3, $4, $5)", role_menu.channel_id.to_string(), role_menu.message_id.to_string(), ctx.guild_id().unwrap().0.to_string(), emoji.to_string(), r.id.to_string()).execute(&pool).await?;
            let all_reactions = sqlx::query_as!(
                ReactionRole,
                "SELECT * FROM reaction_roles WHERE message_id=$1",
                message_id.to_string()
            )
            .fetch_all(&pool)
            .await?;
            let channel = channel_id.to_channel(&ctx.discord()).await?;
            let mut menu_msg = channel
                .guild()
                .unwrap()
                .message(ctx.discord(), role_menu.message_id.parse::<u64>()?)
                .await?;
            
            update_menu(ctx, menu_msg).await?;

            ctx.say("Done! I've added that to the list for you").await?;
        } else {
            ctx.say("Whoops! That role doesn't exist!").await?;
            return Ok(());
        }
    }

    Ok(())
}

/// Removes a reaction from the menu
///
/// Usage:
///     ~rroles del <Message ID> <emoji>
#[poise::command(prefix_command, check = "allowed_to_create_roles")]
pub async fn del(ctx: Context<'_>,
                 #[description = "The Message ID of the menu"]
                 message_id: u64,
                 #[description = "The emoji you want to remove"]
                 emoji: ReactionType,
                 ) -> Result<(), Error> {
    {
        let pool = ctx.data().pg.lock().unwrap().clone();
        let reaction_with_menu = sqlx::query_as!(ReactionRole, "SELECT * FROM reaction_roles WHERE message_id=$1 AND reaction=$2", message_id.to_string(), emoji.to_string()).fetch_one(&pool).await?;
        let channel_id = ChannelId(reaction_with_menu.channel_id.parse::<u64>()?); 
        let channel = channel_id.to_channel(ctx.discord()).await?;
        let mut message = channel.guild().unwrap().message(ctx.discord(), message_id).await?;
        
        // Delete from DB
        // We can just use `ReactionRole.id` here to avoid having to do more complex conditionals
        sqlx::query!("DELETE FROM reaction_roles WHERE id=$1", reaction_with_menu.id).execute(&pool).await?;
        
        message.delete_reaction_emoji(ctx.discord(), emoji).await?;
        
        update_menu(ctx, message).await?;
    }

    ctx.say("Alright! I've removed that emoji from the menu.").await?;

    Ok(())
}

fn get_reactiontype_display(rt: &ReactionType) -> String {
    match rt {
        ReactionType::Unicode(emote) => emote.clone(),
        ReactionType::Custom { id, name, .. } => {
            if let Some(name) = name {
                format!("<:{}:{}>", name, id)
            } else {
                format!("<{}>", id)
            }
        }
        _ => String::new(),
    }
}

fn gen_reaction_list(reacts: Vec<ReactionRole>) -> String {
    let mut rolelist_formatted =
        String::from("Choose the appropriate reaction to gain the role!\n\n");
    for r in reacts {
        rolelist_formatted.push_str(&format!("{} - <@&{}>\n", r.reaction, r.role_id));
    }
    rolelist_formatted
}

async fn update_menu(ctx: Context<'_>, mut msg: serenity::Message) -> Result<serenity::Message, Error> {
    {
        let pool = ctx.data().pg.lock().unwrap().clone();
        let all_reactions = sqlx::query_as!(ReactionRole, "SELECT * FROM reaction_roles WHERE message_id=$1", msg.id.0.to_string()).fetch_all(&pool).await?;
        let rolelist_formatted = gen_reaction_list(all_reactions.clone());
        let title = msg.clone().embeds[0].title.clone().unwrap();
        msg.edit(ctx.discord(), |m| {
            m.embed(|e| {
                e.title(title);
                e.description(rolelist_formatted)
            })
        }).await?;

        for r in all_reactions.iter() {
            msg.react(ctx.discord(), ReactionType::from_str(&r.reaction)?).await?;
        }
    }

    Ok(msg)
}