From 5353b04f29d7d6d82a16c40277771835084a9310 Mon Sep 17 00:00:00 2001 From: Edward Shen Date: Sat, 2 May 2020 23:56:04 -0400 Subject: [PATCH] arknights features --- src/commands/heck.rs | 4 +- src/commands/mod.rs | 5 +- src/commands/op.rs | 381 ++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 2 +- src/util/db.rs | 35 +++- src/util/mod.rs | 1 + src/util/operators.rs | 104 ++++++++++++ 7 files changed, 524 insertions(+), 8 deletions(-) create mode 100644 src/commands/op.rs create mode 100644 src/util/operators.rs diff --git a/src/commands/heck.rs b/src/commands/heck.rs index e8355c4..7355eea 100644 --- a/src/commands/heck.rs +++ b/src/commands/heck.rs @@ -6,9 +6,9 @@ use serenity::prelude::Context; #[command] async fn heck(ctx: &Context, msg: &Message) -> CommandResult { let db_pool = ctx.data.clone(); - let mut db_pool = db_pool.write().await; + let db_pool = db_pool.read().await; let db_pool = db_pool - .get_mut::() + .get::() .expect("No db pool in context?!"); let value = db_pool.get_heck().await; diff --git a/src/commands/mod.rs b/src/commands/mod.rs index aac05a0..0438c71 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -1,6 +1,6 @@ use crate::commands::{ clap::CLAP_COMMAND, crosspost::CROSSPOST_COMMAND, cube::CUBE_COMMAND, heck::HECK_COMMAND, - mock::MOCK_COMMAND, source::SOURCE_COMMAND, + mock::MOCK_COMMAND, op::OP_COMMAND, source::SOURCE_COMMAND, }; use serenity::framework::standard::macros::group; @@ -9,8 +9,9 @@ mod crosspost; mod cube; mod heck; mod mock; +mod op; mod source; #[group] -#[commands(heck, clap, cube, source, crosspost, mock)] +#[commands(heck, clap, cube, source, crosspost, mock, op)] pub(crate) struct General; diff --git a/src/commands/op.rs b/src/commands/op.rs new file mode 100644 index 0000000..4b162c5 --- /dev/null +++ b/src/commands/op.rs @@ -0,0 +1,381 @@ +use crate::{ + util::{ + debug_say, + operators::{get_china_ops, get_global_ops, Operator}, + }, + DbConnPool, +}; +use log::warn; +use serenity::framework::standard::{macros::command, Args, CommandResult}; +use serenity::model::{channel::Message, id::UserId}; +use serenity::{async_trait, prelude::Context}; +use sqlx::Error; +use std::{collections::HashSet, str::FromStr}; + +#[command] +#[sub_commands(list, add, remove, missing, roll, pity)] +async fn op(_: &Context, _: &Message) -> CommandResult { + Ok(()) +} + +#[command] +async fn list(ctx: &Context, msg: &Message, args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + let user_id = match args.current().map(|id| UserId::from_str(id.trim())) { + Some(Ok(user_id)) => user_id, + _ => msg.author.id, + } + .as_u64() + .to_string(); + + debug_say( + msg, + ctx, + format!( + "{}'s 6\u{2605} Operators:\n{}", + msg.author, + db_pool + .get_operators(user_id) + .await? + .iter() + .map(|(op, pot)| format!("{:?} (Pot {})", op, pot)) + .collect::>() + .join("\n") + ), + ) + .await?; + Ok(()) +} + +#[command] +async fn add(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + let mut failed_parses = false; + let mut num_added: usize = 0; + for arg in args.iter::() { + match arg { + Ok(op) => { + db_pool + .add_operator(msg.author.id.as_u64().to_string(), op) + .await?; + num_added += 1; + } + Err(_) => { + failed_parses = true; + } + }; + } + if failed_parses { + debug_say( + msg, + ctx, + "Unable to add some operators. Check your spelling and try again.", + ) + .await?; + } + + match num_added { + 0 => debug_say(msg, ctx, "Didn't add any operators...").await?, + 1 => debug_say(msg, ctx, "Added an operator!").await?, + n => debug_say(msg, ctx, format!("Added {} operators!", n)).await?, + }; + Ok(()) +} + +#[command] +async fn remove(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + let mut failed_parses = false; + let mut num_added: usize = 0; + for arg in args.iter::() { + match arg { + Ok(op) => { + db_pool + .remove_operator(msg.author.id.as_u64().to_string(), op) + .await?; + num_added += 1; + } + Err(_) => { + failed_parses = true; + } + }; + } + if failed_parses { + debug_say( + msg, + ctx, + "Unable to remove some operators. Check your spelling and try again.", + ) + .await?; + } + + match num_added { + 0 => debug_say(msg, ctx, "Didn't remove any operators...").await?, + 1 => debug_say(msg, ctx, "Removed an operator!").await?, + n => debug_say(msg, ctx, format!("Removed {} operators!", n)).await?, + }; + Ok(()) +} + +#[command] +async fn missing(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + let user_id = match args.current().map(|id| UserId::from_str(id.trim())) { + Some(Ok(user_id)) => user_id, + _ => msg.author.id, + } + .as_u64() + .to_string(); + let operators = db_pool + .get_operators(user_id) + .await? + .iter() + .map(|(op, _)| *op) + .collect::>(); + + let compare_ops = match args.single_quoted::() { + Ok(arg) if arg == "china" || arg == "cn" => get_china_ops(), + _ => get_global_ops(), + }; + + let operators = operators + .symmetric_difference(&compare_ops) + .collect::>(); + + let resp = if operators.len() > 5 { + format!( + "Missing {} and {} more...", + operators + .iter() + .take(5) + .map(|op| format!("{:?}", op)) + .collect::>() + .join(", "), + operators.len() - 5, + ) + } else { + format!( + "Missing {}", + operators + .iter() + .map(|op| format!("{:?}", op)) + .collect::>() + .join(", ") + ) + }; + + debug_say(msg, ctx, resp).await?; + + Ok(()) +} + +#[command] +async fn roll(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + if args.current().is_none() { + args = Args::new("1", &[]); + } + + while args.current().is_some() { + match args.quoted().current() { + Some(amt) if amt.parse::().is_ok() => { + db_pool + .add_roll_count(msg.author.id.as_u64().to_string(), amt.parse().unwrap()) + .await?; + } + Some(operator) if Operator::from_str(operator).is_ok() => { + let new_op = Operator::from_str(operator).unwrap(); + db_pool + .add_operator(msg.author.id.as_u64().to_string(), new_op) + .await?; + debug_say(msg, ctx, format!("Congratulations on {:?}!", new_op)).await?; + db_pool + .reset_roll_count(msg.author.id.as_u64().to_string()) + .await?; + } + _ => (), + } + args.advance(); + } + + pity(ctx, msg, args).await?; + + Ok(()) +} + +#[command] +#[sub_commands(reset)] +async fn pity(ctx: &Context, msg: &Message, args: Args) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + let mut other_user = false; + + let user_id = match args.current().map(|id| UserId::from_str(id.trim())) { + Some(Ok(user_id)) => { + other_user = true; + user_id + } + _ => msg.author.id, + } + .as_u64() + .to_string(); + match db_pool.get_roll_count(user_id).await { + Ok(count) => { + debug_say( + msg, + ctx, + format!( + "{}'s Current roll: {}.\n6\u{2605} chance: {}%", + msg.author, + count, + 2 + (count as u16).saturating_sub(50) * 2 + ), + ) + .await?; + } + Err(sqlx::Error::RowNotFound) => { + if !other_user { + reset(ctx, msg, Args::new("", &[])).await?; + } else { + debug_say( + msg, + ctx, + "We don't know where you're currently at. Use ~pity reset first!", + ) + .await?; + } + } + Err(_) => { + warn!("Unable to communicate with database"); + } + } + Ok(()) +} + +#[command] +async fn reset(ctx: &Context, msg: &Message) -> CommandResult { + let db_pool = ctx.data.clone(); + let db_pool = db_pool.read().await; + let db_pool = db_pool + .get::() + .expect("No db pool in context?!"); + + db_pool + .reset_roll_count(msg.author.id.as_u64().to_string()) + .await?; + + pity(ctx, msg, Args::new("", &[])).await?; + Ok(()) +} + +#[async_trait] +trait OpCommandQueries { + async fn get_roll_count(&self, id: String) -> Result; + async fn reset_roll_count(&self, id: String) -> Result<(), Error>; + async fn add_roll_count(&self, id: String, amount: u32) -> Result<(), Error>; + async fn get_operators(&self, id: String) -> Result, Error>; + async fn add_operator(&self, id: String, op: Operator) -> Result<(), Error>; + async fn remove_operator(&self, id: String, op: Operator) -> Result<(), Error>; +} + +#[async_trait] +impl OpCommandQueries for DbConnPool { + async fn get_roll_count(&self, id: String) -> Result { + let resp = sqlx::query!("SELECT count FROM RollCount WHERE user_id = ?", id) + .fetch_one(&self.pool) + .await?; + + Ok(resp.count) + } + + async fn reset_roll_count(&self, id: String) -> Result<(), Error> { + sqlx::query!( + "INSERT INTO RollCount (user_id, count) VALUES (?, 0) + ON CONFLICT(user_id) DO UPDATE SET count = 0", + id + ) + .execute(&self.pool) + .await?; + Ok(()) + } + + async fn add_roll_count(&self, id: String, amount: u32) -> Result<(), Error> { + sqlx::query!( + "INSERT INTO RollCount (user_id, count) VALUES (?, 0) + ON CONFLICT(user_id) DO UPDATE SET count = MIN(99, count + ?)", + id, + amount as i32 + ) + .execute(&self.pool) + .await?; + Ok(()) + } + + async fn get_operators(&self, id: String) -> Result, Error> { + Ok(sqlx::query!( + "SELECT operator, count FROM OperatorCount WHERE user_id = ?", + id + ) + .fetch_all(&self.pool) + .await? + .iter() + .map(|record| { + ( + Operator::from_str(record.operator.as_ref()).unwrap(), + record.count as u32, + ) + }) + .collect()) + } + + async fn add_operator(&self, id: String, op: Operator) -> Result<(), Error> { + sqlx::query!( + "INSERT INTO OperatorCount (user_id, operator, count) VALUES (?, ?, 1) + ON CONFLICT(user_id, operator) DO UPDATE SET count = count + 1", + id, + op + ) + .execute(&self.pool) + .await?; + Ok(()) + } + + async fn remove_operator(&self, id: String, op: Operator) -> Result<(), Error> { + sqlx::query!( + "UPDATE OperatorCount SET count = MAX(count - 1, 0) WHERE user_id = ? AND operator = ?", + id, + op + ) + .execute(&self.pool) + .await?; + + sqlx::query!("DELETE FROM OperatorCount WHERE count = 0") + .execute(&self.pool) + .await?; + + Ok(()) + } +} diff --git a/src/main.rs b/src/main.rs index 30168cf..1bd49fb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,7 @@ mod commands; mod passive; mod util; -pub(crate) const COMMAND_PREFIX: &str = "~"; +pub(crate) const COMMAND_PREFIX: &str = "\\"; #[tokio::main] async fn main() { diff --git a/src/util/db.rs b/src/util/db.rs index cae8339..f63eb9e 100644 --- a/src/util/db.rs +++ b/src/util/db.rs @@ -9,7 +9,7 @@ use std::env; type DbPool = Pool; pub(crate) struct DbConnPool { - pool: DbPool, + pub pool: DbPool, } impl DbConnPool { @@ -40,21 +40,50 @@ async fn init_pool() -> Result { .build(&env::var("DATABASE_URL").unwrap()) .await?; + // Heck table sqlx::query!( - "CREATE TABLE IF NOT EXISTS Heck (id INTEGER PRIMARY KEY NOT NULL, count INTEGER NOT NULL)" + "CREATE TABLE IF NOT EXISTS Heck ( + id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL UNIQUE, + count INTEGER NOT NULL + )" ) .execute(&pool) .await?; debug!("Table Heck exists or was created"); + // Arknights row roll counting + sqlx::query!( + "CREATE TABLE IF NOT EXISTS RollCount ( + user_id TEXT PRIMARY KEY UNIQUE NOT NULL, + count INTEGER NOT NULL + )" + ) + .execute(&pool) + .await?; + + debug!("Table RollCount exists or was created"); + + sqlx::query!( + "CREATE TABLE IF NOT EXISTS OperatorCount ( + user_id TEXT NOT NULL, + operator TEXT NOT NULL, + count INTEGER NOT NULL, + UNIQUE(user_id, operator) + )" + ) + .execute(&pool) + .await?; + + debug!("Table OperatorCount exists or was created"); + if sqlx::query!("SELECT count FROM Heck") .fetch_all(&pool) .await? .is_empty() { debug!("No entries in Heck, inserting default one"); - sqlx::query!("INSERT INTO Heck VALUES (1, 0)") + sqlx::query!("INSERT INTO Heck (count) VALUES (0)") .execute(&pool) .await?; } diff --git a/src/util/mod.rs b/src/util/mod.rs index eaecfa5..4a71c21 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -7,6 +7,7 @@ use serenity::{ use std::env::var; pub mod db; pub mod error; +pub mod operators; lazy_static! { static ref BOT_OWNER_ID: u64 = var("BOT_OWNER_ID") diff --git a/src/util/operators.rs b/src/util/operators.rs new file mode 100644 index 0000000..9ef8ffd --- /dev/null +++ b/src/util/operators.rs @@ -0,0 +1,104 @@ +use log::warn; +use std::{collections::HashSet, str::FromStr}; +#[derive(sqlx::Type, Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum Operator { + Chen, + Siege, + Shining, + Nightingale, + Ifrit, + Eyjafjalla, + Exusiai, + Angelina, + SilverAsh, + Hoshiguma, + Saria, + Skadi, + Schwarz, + Hellagur, + Magallan, + Mostima, + Blaze, + Aak, + Nian, + Ceobe, + Bagpipe, + Phantom, + W, + Weedy, +} + +pub fn get_global_ops() -> HashSet { + let global_ops = &[ + Operator::Chen, + Operator::Siege, + Operator::Shining, + Operator::Nightingale, + Operator::Ifrit, + Operator::Eyjafjalla, + Operator::Exusiai, + Operator::Angelina, + Operator::SilverAsh, + Operator::Hoshiguma, + Operator::Saria, + Operator::Skadi, + Operator::Schwarz, + Operator::Hellagur, + Operator::Magallan, + ]; + let mut set = HashSet::with_capacity(global_ops.len()); + set.extend(global_ops); + set +} + +pub fn get_china_ops() -> HashSet { + let mut global_ops = get_global_ops(); + global_ops.extend(&[ + Operator::Mostima, + Operator::Blaze, + Operator::Aak, + Operator::Nian, + Operator::Ceobe, + Operator::Bagpipe, + Operator::Phantom, + Operator::W, + Operator::Weedy, + ]); + global_ops +} + +impl FromStr for Operator { + type Err = (); + fn from_str(value: &str) -> Result { + match value.to_ascii_lowercase().as_ref() { + "chen" | "ch'en" => Ok(Self::Chen), + "siege" => Ok(Self::Siege), + "shining" => Ok(Self::Shining), + "nightingale" => Ok(Self::Nightingale), + "ifrit" => Ok(Self::Ifrit), + "effy" | "eyja" | "eyjafjalla" => Ok(Self::Eyjafjalla), + "exu" | "exusiai" | "apple pie" | "appuru pai" => Ok(Self::Exusiai), + "angelina" => Ok(Self::Angelina), + "sa" | "silver daddy" | "silverash" => Ok(Self::SilverAsh), + "hoshi" | "hoshiguma" => Ok(Self::Hoshiguma), + "saria" => Ok(Self::Saria), + "skadi" => Ok(Self::Skadi), + "schwarz" => Ok(Self::Schwarz), + "hellagur" => Ok(Self::Hellagur), + "penguin" | "magallan" => Ok(Self::Magallan), + "mostina" => Ok(Self::Mostima), + "blaze" => Ok(Self::Blaze), + "aak" => Ok(Self::Aak), + "nian" => Ok(Self::Nian), + "ceobe" => Ok(Self::Ceobe), + "bagpipe" => Ok(Self::Bagpipe), + "phantom" => Ok(Self::Phantom), + "w" => Ok(Self::W), + "weedy" => Ok(Self::Weedy), + _ => { + warn!("Failed to convert {}", value); + Err(()) + } + } + } +}