arknights features

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::<DbConnPool>()
+        .get::<DbConnPool>()
         .expect("No db pool in context?!");
     let value = db_pool.get_heck().await;
 

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;

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::<DbConnPool>()
+        .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::<Vec<_>>()
+                .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::<DbConnPool>()
+        .expect("No db pool in context?!");
+    let mut failed_parses = false;
+    let mut num_added: usize = 0;
+    for arg in args.iter::<Operator>() {
+        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::<DbConnPool>()
+        .expect("No db pool in context?!");
+    let mut failed_parses = false;
+    let mut num_added: usize = 0;
+    for arg in args.iter::<Operator>() {
+        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::<DbConnPool>()
+        .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::<HashSet<_>>();
+
+    let compare_ops = match args.single_quoted::<String>() {
+        Ok(arg) if arg == "china" || arg == "cn" => get_china_ops(),
+        _ => get_global_ops(),
+    };
+
+    let operators = operators
+        .symmetric_difference(&compare_ops)
+        .collect::<Vec<_>>();
+
+    let resp = if operators.len() > 5 {
+        format!(
+            "Missing {} and {} more...",
+            operators
+                .iter()
+                .take(5)
+                .map(|op| format!("{:?}", op))
+                .collect::<Vec<_>>()
+                .join(", "),
+            operators.len() - 5,
+        )
+    } else {
+        format!(
+            "Missing {}",
+            operators
+                .iter()
+                .map(|op| format!("{:?}", op))
+                .collect::<Vec<_>>()
+                .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::<DbConnPool>()
+        .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::<u64>().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::<DbConnPool>()
+        .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::<DbConnPool>()
+        .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<i32, Error>;
+    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<Vec<(Operator, u32)>, 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<i32, Error> {
+        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<Vec<(Operator, u32)>, 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(())
+    }
+}

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() {

src/util/db.rs

@@ -9,7 +9,7 @@ use std::env;
 type DbPool = Pool<SqliteConnection>;
 
 pub(crate) struct DbConnPool {
-    pool: DbPool,
+    pub pool: DbPool,
 }
 
 impl DbConnPool {
@@ -40,21 +40,50 @@ async fn init_pool() -> Result<DbPool, Error> {
         .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?;
     }

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")

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<Operator> {
+    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<Operator> {
+    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<Self, Self::Err> {
+        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(())
+            }
+        }
+    }
+}