From 3d0dc324a6865a3237a3cfc7fae23c44d95bfade Mon Sep 17 00:00:00 2001
From: Owen <owenrees82@gmail.com>
Date: Mon, 13 Jan 2025 15:37:32 +0100
Subject: [PATCH] custom rate limits

---
 .../cogs/plantnet_message.ex                  |  4 ++--
 lib/plantid_discord_bot/rate_limiter.ex       | 22 ++++++++++++++-----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/lib/plantid_discord_bot/cogs/plantnet_message.ex b/lib/plantid_discord_bot/cogs/plantnet_message.ex
index 9dfbe6d..766eeda 100644
--- a/lib/plantid_discord_bot/cogs/plantnet_message.ex
+++ b/lib/plantid_discord_bot/cogs/plantnet_message.ex
@@ -17,14 +17,14 @@ defmodule PlantIdDiscordBot.Cog.PlantNetMessage do
     IO.inspect(message)
 
     case RateLimiter.check_limit(message.guild_id) do
-      {:limit_exceeded, _} ->
+      {:limit_exceeded, _requests_used, _requests_limit} ->
         Api.create_message(message.channel_id,
           content:
             "This server has exceeded its allowed requests in 24 hours. Please try again tomorrow.",
           message_reference: %{message_id: message.id}
         )
 
-      {:ok, _} ->
+      {:ok, _requests_used, _requests_limit} ->
         do_identification(message)
     end
   end
diff --git a/lib/plantid_discord_bot/rate_limiter.ex b/lib/plantid_discord_bot/rate_limiter.ex
index e864bd5..2bfd83b 100644
--- a/lib/plantid_discord_bot/rate_limiter.ex
+++ b/lib/plantid_discord_bot/rate_limiter.ex
@@ -2,7 +2,10 @@ defmodule PlantIdDiscordBot.RateLimiter do
   use GenServer
   require Logger
 
-  @request_limit Application.compile_env(:plantid_discord_bot, :guild_request_limit_per_day)
+  @default_request_limit Application.compile_env(
+                           :plantid_discord_bot,
+                           :guild_request_limit_per_day
+                         )
 
   @custom_limits %{
     1_002_507_312_159_797_318 => 100,
@@ -34,19 +37,26 @@ defmodule PlantIdDiscordBot.RateLimiter do
   @doc """
   Check if the number of requests for a guild exceeds the limit.
   """
-  @spec check_limit(String.t()) :: {:ok, integer()} | {:limit_exceeded, integer()}
+  @spec check_limit(String.t() | integer()) :: {:ok, integer()} | {:limit_exceeded, integer()}
   def check_limit(guild_id) do
+    request_limit =
+      case guild_id in Map.keys(@custom_limits) do
+        true -> @custom_limits[guild_id]
+        _ -> @default_request_limit
+      end
+
     case get(guild_id) do
       {:ok, value} ->
-        if value >= @request_limit do
-          {:limit_exceeded, value}
+        if value >= request_limit do
+          {:limit_exceeded, value, request_limit}
         else
-          {:ok, value}
+          {:ok, value, request_limit}
         end
 
       {:error, _} ->
         put(guild_id, 0)
-        {:ok, 0}
+        Logger.debug("Set counter for guild #{guild_id} to 0")
+        {:ok, 0, request_limit}
     end
   end