diff --git a/src/commands/utilities/ai.ts b/src/commands/utilities/ai.ts index c07227a..5f4893f 100644 --- a/src/commands/utilities/ai.ts +++ b/src/commands/utilities/ai.ts @@ -390,18 +390,39 @@ async function incrementAndCheckDailyLimit(userId: string, limit = 20): Promise< const client = await pool.connect(); try { await client.query('BEGIN'); + + const voteCheck = await client.query( + `SELECT vote_timestamp FROM votes + WHERE user_id = $1 + AND vote_timestamp > NOW() - INTERVAL '24 hours' + ORDER BY vote_timestamp DESC + LIMIT 1`, + [userId] + ); + + const hasVotedRecently = voteCheck.rows.length > 0; + const effectiveLimit = hasVotedRecently ? limit + 10 : limit; + await client.query('INSERT INTO users (user_id) VALUES ($1) ON CONFLICT (user_id) DO NOTHING', [ userId, ]); + const res = await client.query( - `INSERT INTO ai_usage (user_id, usage_date, count) VALUES ($1, $2, 1) - ON CONFLICT (user_id, usage_date) DO UPDATE SET count = ai_usage.count + 1 RETURNING count`, + `INSERT INTO ai_usage (user_id, usage_date, count) + VALUES ($1, $2, 1) + ON CONFLICT (user_id, usage_date) + DO UPDATE SET count = ai_usage.count + 1 + RETURNING count`, [userId, today], ); + await client.query('COMMIT'); - return res.rows[0].count <= limit; + + return res.rows[0].count <= effectiveLimit; + } catch (err) { await client.query('ROLLBACK'); + console.error('Error in incrementAndCheckDailyLimit:', err); throw err; } finally { client.release(); diff --git a/src/utils/topgg.ts b/src/utils/topgg.ts index ee9e468..01fac0c 100644 --- a/src/utils/topgg.ts +++ b/src/utils/topgg.ts @@ -6,7 +6,7 @@ export async function checkVoteStatus( const now = new Date(); const nextVote = new Date(now.getTime() + VOTE_COOLDOWN_HOURS * 60 * 60 * 1000); - console.log(`Vote recorded for user ${userId}. Next vote available at ${nextVote.toISOString()}`); + console.log(`Vote check for user ${userId}. Next vote available at ${nextVote.toISOString()}`); return { hasVoted: true, diff --git a/src/utils/voteManager.ts b/src/utils/voteManager.ts index efeb44e..4776230 100644 --- a/src/utils/voteManager.ts +++ b/src/utils/voteManager.ts @@ -150,14 +150,16 @@ export async function recordVote( await client.query( `INSERT INTO message_credits (user_id, credits_remaining, last_reset) - VALUES ($1, $2, NOW()) - ON CONFLICT (user_id) WHERE server_id IS NULL - DO UPDATE SET - credits_remaining = message_credits.credits_remaining + $2, - last_reset = NOW() - RETURNING credits_remaining`, + VALUES ($1, $2, NOW()) + ON CONFLICT (user_id) + DO UPDATE SET + credits_remaining = message_credits.credits_remaining + EXCLUDED.credits_remaining, + last_reset = NOW() + RETURNING credits_remaining`, [userId, VOTE_CREDITS], ); + + console.log(`Added ${VOTE_CREDITS} credits to user ${userId} (global)`); const clientBot = new Client({ intents: [GatewayIntentBits.Guilds, GatewayIntentBits.GuildMembers], @@ -179,13 +181,14 @@ export async function recordVote( await client.query( `INSERT INTO message_credits (user_id, server_id, credits_remaining, last_reset) VALUES ($1, $2, $3, NOW()) - ON CONFLICT (user_id, server_id) WHERE server_id IS NOT NULL + ON CONFLICT (user_id, server_id) DO UPDATE SET credits_remaining = message_credits.credits_remaining + $3, last_reset = NOW() RETURNING credits_remaining`, [userId, guild.id, VOTE_CREDITS], ); + console.log(`Added ${VOTE_CREDITS} credits to user ${userId} in server ${guild.id}`); } } catch (error) { console.error(`Error processing guild ${guild.id}:`, error); @@ -199,6 +202,44 @@ export async function recordVote( clientBot.destroy().catch(console.error); } + await client.query( + `INSERT INTO ai_usage (user_id, usage_date, count) + VALUES ($1, CURRENT_DATE, 10) + ON CONFLICT (user_id, usage_date) + DO UPDATE SET + count = GREATEST(0, ai_usage.count) + 10 + RETURNING count`, + [userId] + ); + + console.log(`Added 10 AI usage credits to user ${userId}`); + + try { + const clientBot = new Client({ + intents: [GatewayIntentBits.Guilds, GatewayIntentBits.GuildMembers, GatewayIntentBits.MessageContent], + }); + + await clientBot.login(process.env.TOKEN); + const user = await clientBot.users.fetch(userId); + + if (user) { + const nextVoteTime = Math.floor((Date.now() + 12 * 60 * 60 * 1000) / 1000); + await user.send( + `🎉 **Thank you for voting for Aethel!**\n` + + `\n` + + `You've received **+10 AI message credits** for today!\n` + + `\n` + + `You can vote again \n` + + `\n` + + `Thank you for your support! ❤️` + ).catch(console.error); + } + + clientBot.destroy().catch(console.error); + } catch (error) { + console.error('Failed to send vote thank you DM:', error); + } + await client.query('COMMIT'); return { @@ -242,29 +283,47 @@ export async function getRemainingCredits(userId: string, serverId?: string): Pr export async function canUseAIFeature( userId: string, - serverId?: string, + _serverId?: string, // Prefix with underscore to indicate intentionally unused ): Promise<{ canUse: boolean; remainingCredits: number }> { + const client = await pool.connect(); try { - const credits = await getRemainingCredits(userId, serverId); - if (credits.remaining > 0) { - const updateQuery = serverId - ? 'UPDATE message_credits SET credits_remaining = credits_remaining - 1 WHERE user_id = $1 AND server_id = $2 RETURNING credits_remaining' - : 'UPDATE message_credits SET credits_remaining = credits_remaining - 1 WHERE user_id = $1 AND server_id IS NULL RETURNING credits_remaining'; - - const result = await pool.query(updateQuery, [userId, serverId].filter(Boolean)); - + await client.query('BEGIN'); + + const result = await client.query( + `SELECT count FROM ai_usage + WHERE user_id = $1 AND usage_date = CURRENT_DATE + FOR UPDATE`, + [userId] + ); + + if (result.rows.length > 0 && result.rows[0].count > 0) { + const updateResult = await client.query( + `UPDATE ai_usage + SET count = count - 1 + WHERE user_id = $1 AND usage_date = CURRENT_DATE + RETURNING count`, + [userId] + ); + + await client.query('COMMIT'); + return { canUse: true, - remainingCredits: result.rows[0]?.credits_remaining || 0, + remainingCredits: updateResult.rows[0]?.count || 0, }; } - + + await client.query('COMMIT'); + return { canUse: false, remainingCredits: 0, }; } catch (error) { + await client.query('ROLLBACK'); console.error('Error checking AI feature usage:', error); throw new Error('Failed to check AI feature usage'); + } finally { + client.release(); } }