From 6792c059597c37accaf70ea584f53c50165fdc68 Mon Sep 17 00:00:00 2001
From: Wroclaw <wroclaw223@outlook.com>
Date: Thu, 21 Sep 2023 20:07:35 +0200
Subject: [PATCH] Quota: add tokenCount QuotaEngine

---
 src/quota/messageCount.ts |   2 +-
 src/quota/tokenCount.ts   | 124 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+), 1 deletion(-)
 create mode 100644 src/quota/tokenCount.ts

diff --git a/src/quota/messageCount.ts b/src/quota/messageCount.ts
index a38cfdb..7b97755 100644
--- a/src/quota/messageCount.ts
+++ b/src/quota/messageCount.ts
@@ -101,7 +101,7 @@ export default class MessageCount implements IQuota {
   }
 }
 
-function milisecondsToHumanReadable(totalMiliseconds: number): string {
+export function milisecondsToHumanReadable(totalMiliseconds: number): string {
   const negative = totalMiliseconds < 0;
   if (negative) totalMiliseconds = -totalMiliseconds;
 
diff --git a/src/quota/tokenCount.ts b/src/quota/tokenCount.ts
new file mode 100644
index 0000000..f1ce0a3
--- /dev/null
+++ b/src/quota/tokenCount.ts
@@ -0,0 +1,124 @@
+import { User, GuildMember } from "discord.js";
+
+import IQuota, { userQuotaData, userQuotaRecoveryData } from "../IQuota";
+import { apiRequest } from "../execution";
+import { database } from "../index";
+import { milisecondsToHumanReadable } from "./messageCount";
+import { Usage } from "@prisma/client";
+
+/**
+ * Quota based on Tokens used and generated
+ */
+export default class tokenCount implements IQuota {
+  defaultQuota: number;
+  lookback: number;
+  considerInputTokensAsHalf: boolean;
+
+  constructor(
+    defaultQuota: number = 512 * 25,
+    lookback: number = 1000 * 60 * 60 * 24,
+    considerInputTokensAsHalf: boolean = true,
+  ) {
+    this.defaultQuota = defaultQuota;
+    this.lookback = lookback;
+    this.considerInputTokensAsHalf = considerInputTokensAsHalf;
+  }
+
+  private getUserQuota(id: string) {
+    return database.limits.findUnique({
+      where: { user: BigInt(id) },
+    });
+  }
+
+  async checkUser(
+    user: string | User | GuildMember,
+    request: apiRequest
+  ): Promise<userQuotaData> {
+    const userId: string = typeof user === "string" ? user : user.id;
+
+    const userQuota = await this.getUserQuota(userId);
+
+    const usedTokens = (await database.usage.aggregate({
+      _sum: {
+        usageRequest: true,
+        usageResponse: true,
+      },
+      where: {
+        user: BigInt(userId),
+        timestamp: {
+          gte: new Date(request.createdAt.getTime() - this.lookback),
+        }
+      }
+    }))._sum;
+
+    if (!usedTokens.usageRequest || !usedTokens.usageResponse) throw new Error("Null from a database!! (tokenCount Quota)");
+
+    const usedUnits = (() => {
+      if (this.considerInputTokensAsHalf)
+        return usedTokens.usageResponse + usedTokens.usageRequest / 2;
+      return usedTokens.usageResponse + usedTokens.usageRequest;
+    })();
+
+    if (userQuota?.vip) return this.createUserQuotaData(Infinity, usedUnits);
+    if (userQuota?.limit)
+      return this.createUserQuotaData(userQuota.limit, usedUnits);
+    return this.createUserQuotaData(this.defaultQuota, usedUnits);
+  }
+
+  findNthUsage(user: string, requestTimestamp: number, unitCount: number) {
+    if (this.considerInputTokensAsHalf)
+      throw("Not implemented");
+    return database.$queryRaw<Array<Usage & {usage: bigint}>>`
+      SELECT t1.*, (
+          SELECT
+            SUM(usageResponse + usageRequest) AS usage
+          FROM \`usage\`
+          WHERE 
+            user = ${user} AND
+            timestamp >= ${requestTimestamp - this.lookback} AND
+            timestamp <= t1.timestamp
+        ) as usage
+      FROM
+        \`usage\` AS t1
+      WHERE 
+        user = ${user} AND
+        timestamp >= ${requestTimestamp - this.lookback} AND
+        usage >= ${unitCount}
+      ORDER BY timestamp ASC
+      LIMIT 1
+    `;
+  }
+  async getUserQuotaRecovery(
+    user: string | User | GuildMember,
+    request: apiRequest,
+    unitCount: number = 1
+  ): Promise<userQuotaRecoveryData> {
+    const userId = typeof user ==="string" ? user : user.id;
+
+    const [userQuota, renameMebecause] = await Promise.all([
+      this.checkUser(userId, request),
+      this.findNthUsage(userId, request.createdTimestamp, unitCount)
+    ]);
+
+    console.log(renameMebecause);
+    return {
+      ...userQuota,
+      recoveryTimestamp: (renameMebecause.at(0)?.timestamp.valueOf() ?? Infinity) + this.lookback,
+    };
+  }
+
+  /**
+   * helper funtion to create userQuotaData
+   */
+  private createUserQuotaData(quota: number, used: number): userQuotaData {
+    const humanReadable = milisecondsToHumanReadable(this.lookback);
+    return {
+      quota: quota,
+      used: used,
+      unitName: "tokens",
+      toString() {
+        return `${this.used} of ${this.quota} ${this.unitName} in ${humanReadable}`;
+      },
+    };
+  }
+}