Merge pull request #4 from minneelyyyy/dev

Dev

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 {
-    let total_seconds = duration.as_secs();
-    let seconds = total_seconds % 60;
-    let minutes = (total_seconds / 60) % 60;
-    let hours = total_seconds / 3600;
+use std::time::Duration;
 
-    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 = *data.dailies.read().await.get(&user).unwrap_or(&day_ago);
+    let last = get_last(&mut *tx, user).await?;
+    let existed = last.is_some();
+    let last = last.unwrap_or(Local.timestamp_opt(0, 0).unwrap());
 
-    if last <= day_ago {
-        data.dailies.write().await.insert(user, Instant::now());
+    let now = Local::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;
-        let mut tx = db.begin().await?;
+    if now > next_daily {
+        let mut begin = "".to_string();
+        let mut end = "".to_string();
 
-        let bal = super::get_balance(user, &mut *tx).await?;
-        super::change_balance(user, bal + 50, &mut *tx).await?;
+        let streak = if now < time_to_redeem {
+            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 will be available in **{}**.", format_duration(next))).await?;
+        ctx.reply(format!("Your next daily is not available! It will be available on {}.", next_daily.to_rfc2822())).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 tokio::sync::{Mutex, RwLock};
-use std::collections::HashMap;
-use poise::{serenity_prelude::UserId, ReplyHandle};
+use poise::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 {
-                    database,
-                    mentions: Arc::new(Mutex::new(HashMap::new())),
-                    dailies: Arc::new(RwLock::new(HashMap::new())),
-                })
+                Ok(Data { database })
             })
         })
         .build();