add daily streaks

Cargo.lock

@@ -2190,6 +2190,7 @@ dependencies = [
  "atoi",
  "byteorder",
  "bytes",
+ "chrono",
  "crc",
  "crossbeam-queue",
  "either",
@@ -2271,6 +2272,7 @@ dependencies = [
  "bitflags 2.6.0",
  "byteorder",
  "bytes",
+ "chrono",
  "crc",
  "digest",
  "dotenvy",
@@ -2312,6 +2314,7 @@ dependencies = [
  "base64 0.22.1",
  "bitflags 2.6.0",
  "byteorder",
+ "chrono",
  "crc",
  "dotenvy",
  "etcetera",
@@ -2347,6 +2350,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d5b2cf34a45953bfd3daaf3db0f7a7878ab9b7a6b91b422d24a7a9e4c857b680"
 dependencies = [
  "atoi",
+ "chrono",
  "flume",
  "futures-channel",
  "futures-core",

Cargo.toml

@@ -11,7 +11,7 @@ rand = "0.8.5"
 clap = { version = "4.5.20", features = ["derive"] }
 lamm = { git = "https://github.com/minneelyyyy/lamm", branch = "dev" }
 regex = "1.11.1"
-sqlx = { version = "0.8", features = [ "runtime-tokio-native-tls", "postgres" ] }
+sqlx = { version = "0.8", features = [ "runtime-tokio-native-tls", "postgres", "chrono" ] }
 hex_color = "3"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0", features = ["raw_value"] }

src/commands/gambling/daily.rs

@@ -1,41 +1,132 @@
 use crate::{Context, Error};
 
-use std::time::{Duration, Instant};
+use poise::serenity_prelude::UserId;
+use sqlx::{types::chrono::{DateTime, Local, TimeZone}, PgExecutor, Row};
 
-fn format_duration(duration: Duration) -> String {
+use std::time::Duration;
-    let total_seconds = duration.as_secs();
-    let seconds = total_seconds % 60;
-    let minutes = (total_seconds / 60) % 60;
-    let hours = total_seconds / 3600;
 
-    format!("{:02}:{:02}:{:02}", hours, minutes, seconds)
+async fn get_streak<'a, E>(db: E, user: UserId) -> Result<Option<i32>, Error>
+where
+    E: PgExecutor<'a>,
+{
+    match sqlx::query(
+            "SELECT streak FROM dailies WHERE userid = $1"
+        ).bind(user.get() as i64).fetch_one(db).await
+    {
+        Ok(row) => Ok(Some(row.get(0))),
+        Err(sqlx::Error::RowNotFound) => Ok(None),
+        Err(e) => Err(Box::new(e)),
+    }
+}
+
+async fn set_streak<'a, E>(db: E, user: UserId, streak: i32) -> Result<(), Error>
+where
+    E: PgExecutor<'a>,
+{
+    sqlx::query("INSERT INTO dailies (userid, streak) VALUES ($1, $2) ON CONFLICT (userid) DO UPDATE SET streak = EXCLUDED.streak")
+        .bind(user.get() as i64)
+        .bind(streak)
+        .execute(db).await?;
+
+    Ok(())
+}
+
+async fn get_last<'a, E>(db: E, user: UserId) -> Result<Option<DateTime<Local>>, Error>
+where
+    E: PgExecutor<'a>,
+{
+    match sqlx::query(
+            "SELECT last FROM dailies WHERE userid = $1"
+        ).bind(user.get() as i64).fetch_one(db).await
+    {
+        Ok(row) => Ok(Some(row.get(0))),
+        Err(sqlx::Error::RowNotFound) => Ok(None),
+        Err(e) => Err(Box::new(e)),
+    }
+}
+
+async fn set_last<'a, E>(db: E, user: UserId, last: DateTime<Local>) -> Result<(), Error>
+where
+    E: PgExecutor<'a>,
+{
+    sqlx::query("INSERT INTO dailies (userid, last) VALUES ($1, $2) ON CONFLICT (userid) DO UPDATE SET last = EXCLUDED.last")
+        .bind(user.get() as i64)
+        .bind(last)
+        .execute(db).await?;
+
+    Ok(())
 }
 
-/// Redeem 50 daily tokens.
 #[poise::command(slash_command, prefix_command)]
-pub async fn daily(ctx: Context<'_>) -> Result<(), Error> {
+pub async fn streak(ctx: Context<'_>) -> Result<(), Error> {
+    let db = &ctx.data().database;
+
+    ctx.reply(format!("You have a daily streak of **{}**", get_streak(db, ctx.author().id).await?.unwrap_or(0))).await?;
+    Ok(())
+}
+
+async fn do_claim(ctx: Context<'_>) -> Result<(), Error> {
     let data = ctx.data();
     let user = ctx.author().id;
+    let mut tx = data.database.begin().await?;
 
-    let day_ago = Instant::now() - Duration::from_secs(24 * 60 * 60);
+    let last = get_last(&mut *tx, user).await?;
-    let last = *data.dailies.read().await.get(&user).unwrap_or(&day_ago);
+    let existed = last.is_some();
+    let last = last.unwrap_or(Local.timestamp_opt(0, 0).unwrap());
 
-    if last <= day_ago {
+    let now = Local::now();
-        data.dailies.write().await.insert(user, Instant::now());
+    let next_daily = last + Duration::from_secs(24 * 60 * 60);
+    let time_to_redeem = next_daily + Duration::from_secs(24 * 60 * 60);
 
-        let db = &data.database;
+    if now > next_daily {
-        let mut tx = db.begin().await?;
+        let mut begin = "".to_string();
+        let mut end = "".to_string();
 
-        let bal = super::get_balance(user, &mut *tx).await?;
+        let streak = if now < time_to_redeem {
-        super::change_balance(user, bal + 50, &mut *tx).await?;
+            let streak = get_streak(&mut *tx, user).await?.unwrap_or(0);
+
+            if existed {
+                begin = format!("You have a streak of **{streak}**! ");
+            }
+
+            streak
+        } else {
+            if existed {
+                begin = "You have not redeemed your daily in time and your streak has been reset. ".to_string();
+            }
+
+            0
+        };
+
+        if !existed {
+            end = " Keep redeeming your daily to build up a streak of up to 7 days!".to_string();
+        }
+
+        let payout = 50 + 10 * streak.min(7);
+        
+        let balance = super::get_balance(user, &mut *tx).await?;
+        super::change_balance(user, balance + payout, &mut *tx).await?;
+
+        set_streak(&mut *tx, user, streak + 1).await?;
+        set_last(&mut *tx, user, now).await?;
 
         tx.commit().await?;
 
-        ctx.reply(format!("**50** tokens have been added to your balance.")).await?;
+        ctx.reply(format!("{begin}**{payout}** tokens were added to your balance.{end}")).await?;
     } else {
-        let next = Duration::from_secs(24 * 60 * 60) - last.elapsed();
+        ctx.reply(format!("Your next daily is not available! It will be available on {}.", next_daily.to_rfc2822())).await?;
-        ctx.reply(format!("Your next daily will be available in **{}**.", format_duration(next))).await?;
     }
 
     Ok(())
 }
+
+#[poise::command(slash_command, prefix_command)]
+pub async fn claim(ctx: Context<'_>) -> Result<(), Error> {
+    do_claim(ctx).await
+}
+
+/// Redeem daily tokens.
+#[poise::command(slash_command, prefix_command, subcommands("streak", "claim"))]
+pub async fn daily(ctx: Context<'_>) -> Result<(), Error> {
+    do_claim(ctx).await
+}

src/common.rs

@@ -1,15 +1,8 @@
-use std::sync::Arc;
+use poise::ReplyHandle;
-use tokio::sync::{Mutex, RwLock};
-use std::collections::HashMap;
-use poise::{serenity_prelude::UserId, ReplyHandle};
 use sqlx::{Pool, Postgres};
 
 pub struct Data {
     pub database: Pool<Postgres>,
-    pub mentions: Arc<Mutex<HashMap<UserId, std::time::Instant>>>,
-
-    /// last time the user redeemed a daily
-    pub dailies: Arc<RwLock<HashMap<UserId, std::time::Instant>>>,
 }
 
 pub type Error = Box<dyn std::error::Error + Send + Sync>;

src/main.rs

@@ -4,13 +4,10 @@ pub mod common;
 pub mod inventory;
 use crate::common::{Context, Error, Data};
 
-use std::collections::HashMap;
 use std::env;
 use std::sync::Arc;
 
 use poise::serenity_prelude::{self as serenity};
-use tokio::sync::{Mutex, RwLock};
-
 use clap::Parser;
 
 use sqlx::postgres::PgPoolOptions;
@@ -112,13 +109,19 @@ async fn main() -> Result<(), Error> {
                     "#
                 ).execute(&database).await?;
 
+                sqlx::query(
+                    r#"
+                    CREATE TABLE IF NOT EXISTS dailies (
+                        userid BIGINT NOT NULL PRIMARY KEY,
+                        last TIMESTAMPTZ,
+                        streak INT
+                    )
+                    "#
+                ).execute(&database).await?;
+
                 println!("Bot is ready!");
 
-                Ok(Data {
+                Ok(Data { database })
-                    database,
-                    mentions: Arc::new(Mutex::new(HashMap::new())),
-                    dailies: Arc::new(RwLock::new(HashMap::new())),
-                })
             })
         })
         .build();