rroles: Implement Reaction Roles

now requires postgresql.

5 files changed, 213 insertions, 3 deletions

diff --git a/src/commands/mod.rs b/src/commands/mod.rs
index 50636a3..2121532 100644
--- a/src/commands/mod.rs
+++ b/src/commands/mod.rs
@@ -2,3 +2,4 @@ pub mod meta;
 pub mod pony;
 pub mod actions;
 pub mod osu;
+pub mod reactionroles;
diff --git a/src/commands/reactionroles.rs b/src/commands/reactionroles.rs
new file mode 100644
index 0000000..53aa794
--- /dev/null
+++ b/src/commands/reactionroles.rs
@@ -0,0 +1,146 @@
+use std::{time::Duration, str::FromStr};
+
+use crate::{Context, Error};
+use poise::serenity_prelude::{self as serenity, ReactionType, Emoji, ArgumentConvert};
+use ::serenity::framework::standard::{Args, Delimiter};
+
+#[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)?;
+            let member_permissions = member.permissions().await?;
+
+            Ok(member_permissions.manage_roles())
+    }
+    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! Reply to this message with a list of role names and reactions, like this:\n\n:female_sign::She/Her,:male_sign::He/Him", 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();
+        if let Some(roles) = ctx.author().clone().await_reply(ctx.discord()).timeout(Duration::from_secs(30)).await {
+            let mut args = Args::new(&roles.content, &[Delimiter::Single(',')]);
+
+            let mut rolelist_formatted = String::from("Choose the appropriate reaction to gain the role!\n");
+
+            for a in args.iter::<String>() {
+                if let Ok(tuple) = a {
+                    let split_str: Vec<&str> = tuple.split(':').collect();
+                    let reaction = ReactionType::from_str(split_str[0])?;
+                    let role_name = split_str[1];
+                    if let Some(role) = ctx.guild().unwrap().role_by_name(role_name) {
+                        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?;
+                        rolelist_formatted.push_str(&format!("{} - {}\n", get_reactiontype_display(&reaction), role.name.clone()));
+                    } else {
+                        ctx.say(format!("Invalid role provided: {}", role_name)).await?;
+                        return Ok(());
+                    }
+                }
+            }
+            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(())
+}
+
+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(),
+    }
+}
diff --git a/src/handler.rs b/src/handler.rs
new file mode 100644
index 0000000..933bd72
--- /dev/null
+++ b/src/handler.rs
@@ -0,0 +1,36 @@
+use poise::serenity_prelude as serenity;
+
+use crate::{Data, Error};
+use crate::models::ReactionRole;
+
+pub async fn event_handler(ctx: &serenity::Context,
+                           event: &poise::Event<'_>,
+                           data: &Data,
+                           ) -> Result<(), Error> {
+    {
+        let pool = data.pg.lock().unwrap().clone();
+        match event {
+            poise::Event::ReactionAdd { add_reaction } => {
+                let rrole = sqlx::query_as!(ReactionRole, "SELECT * FROM reaction_roles WHERE message_id=$1 AND reaction=$2", add_reaction.message_id.0.to_string(), add_reaction.emoji.to_string()).fetch_one(&pool).await?;
+                let member = ctx.http.get_member(rrole.guild_id.parse::<u64>()?, add_reaction.user_id.unwrap().0).await?;
+                let member_roles = member.roles;
+                let role_id = serenity::RoleId(rrole.role_id.parse::<u64>()?);
+                if member_roles.contains(&role_id) {
+                    ctx.http.remove_member_role(member.guild_id.0, member.user.id.0, role_id.0, Some("Reaction Role Menu")).await?; 
+                } else {
+                    ctx.http.add_member_role(rrole.guild_id.parse::<u64>()?, add_reaction.user_id.unwrap().0, rrole.role_id.parse::<u64>()?, Some("Reaction Role")).await?;
+                }
+
+                let dm_chan = add_reaction.user_id.unwrap().create_dm_channel(&ctx.http).await?;
+                dm_chan.say(ctx, format!("Toggled the role!")).await?;
+
+                add_reaction.delete(&ctx.http).await?;
+
+            },
+            _ => (),
+        }
+    }
+
+
+    Ok(())
+}
diff --git a/src/main.rs b/src/main.rs
index bb686d2..021886f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,14 +1,17 @@
-use std::time::Duration;
+use std::{time::Duration, sync::Mutex};
 
 use dotenv::dotenv;
+use sqlx::{PgPool, postgres::PgPoolOptions};
 
 type Error = Box<dyn std::error::Error + Send + Sync>;
 type Context<'a> = poise::Context<'a, Data, Error>;
 
 mod commands;
+mod handler;
+mod models;
 
 pub struct Data {
-
+   pg: Mutex<PgPool>, 
 }
 
 /// Show help menu
@@ -63,6 +66,13 @@ async fn main() {
 
             commands::osu::osup(),
             commands::osu::osubm(),
+
+            poise::Command {
+                subcommands: vec![
+                    commands::reactionroles::init(),
+                ],
+                ..commands::reactionroles::rroles()
+            },
         ],
         on_error: |error| Box::pin(on_error(error)),
         pre_command: |ctx| {
@@ -85,6 +95,7 @@ async fn main() {
             ], 
             ..Default::default()
         },
+        listener: |ctx, event, _, data| Box::pin(handler::event_handler(ctx, event, data)),
         ..Default::default()
     };
 
@@ -92,7 +103,14 @@ async fn main() {
         .token(std::env::var("DISCORD_TOKEN").unwrap_or("BAD-TOKEN".into()))
         .user_data_setup(move |_ctx, _ready, _framework| {
             Box::pin(async move {
-                Ok(Data {})
+                let pool = PgPoolOptions::new()
+                    .max_connections(5)
+                    .connect(&std::env::var("DATABASE_URL").unwrap_or("postgres://postgres@localhost/glitch".to_string()))
+                    .await
+                    .expect("Couldn't connect to postgresql");
+                Ok(Data {
+                    pg: Mutex::new(pool)
+                })
             })
         })
         .options(options)
diff --git a/src/models.rs b/src/models.rs
new file mode 100644
index 0000000..4b2e7bb
--- /dev/null
+++ b/src/models.rs
@@ -0,0 +1,9 @@
+#[derive(Debug)]
+pub struct ReactionRole {
+    pub id: i32,
+    pub channel_id: String,
+    pub message_id: String,
+    pub guild_id: String,
+    pub reaction: String,
+    pub role_id: String,
+}