diff --git a/Cargo.lock b/Cargo.lock index 5b35a14..259a30a 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index dbfd274..8bc6474 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/src/commands/gambling/daily.rs b/src/commands/gambling/daily.rs index 6217ca4..8b47246 100644 --- a/src/commands/gambling/daily.rs +++ b/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, 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>, 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) -> 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 +} diff --git a/src/common.rs b/src/common.rs index 7f892f3..dc09de2 100644 --- a/src/common.rs +++ b/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, - pub mentions: Arc>>, - - /// last time the user redeemed a daily - pub dailies: Arc>>, } pub type Error = Box; diff --git a/src/main.rs b/src/main.rs index b215299..22b00d4 100644 --- a/src/main.rs +++ b/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();