From 36be588d7b8c74d64df76a39f6f822855ed85c4a Mon Sep 17 00:00:00 2001
From: Ethan <47520067+Skelmis@users.noreply.github.com>
Date: Thu, 4 Jan 2024 17:45:15 +1300
Subject: [PATCH] Feat/proxy (#71)

* Revert "chore: version 3.19 (#65)" (#69)

This reverts commit 222da50e8b43207a1c646651953854c4f2bd9772.

* feat: add support for proxy
---
 suggestions/bot.py                 | 15 ++++++++++++++-
 suggestions/cogs/help_guild_cog.py | 22 ++++++----------------
 suggestions/main.py                |  2 +-
 suggestions/zonis_routes.py        | 25 ++++++-------------------
 4 files changed, 27 insertions(+), 37 deletions(-)

diff --git a/suggestions/bot.py b/suggestions/bot.py
index 0c814d9..347f075 100644
--- a/suggestions/bot.py
+++ b/suggestions/bot.py
@@ -46,7 +46,7 @@
 
 class SuggestionsBot(commands.AutoShardedInteractionBot, BotBase):
     def __init__(self, *args, **kwargs):
-        self.version: str = "Public Release 3.20"
+        self.version: str = "Public Release 3.19"
         self.main_guild_id: int = 601219766258106399
         self.legacy_beta_role_id: int = 995588041991274547
         self.automated_beta_role_id: int = 998173237282361425
@@ -110,6 +110,19 @@ def __init__(self, *args, **kwargs):
 
         self.zonis: ZonisRoutes = ZonisRoutes(self)
 
+    async def launch_shard(
+        self, _gateway: str, shard_id: int, *, initial: bool = False
+    ) -> None:
+        # Use the proxy if set, else fall back to whatever is default
+        proxy: Optional[str] = os.environ.get("GW_PROXY", _gateway)
+        return await super().launch_shard(proxy, shard_id, initial=initial)
+
+    async def before_identify_hook(
+        self, _shard_id: int | None, *, initial: bool = False  # noqa: ARG002
+    ) -> None:
+        # gateway-proxy
+        return
+
     async def get_or_fetch_channel(self, channel_id: int) -> WrappedChannel:
         try:
             return await super().get_or_fetch_channel(channel_id)
diff --git a/suggestions/cogs/help_guild_cog.py b/suggestions/cogs/help_guild_cog.py
index 55f210e..4b1d602 100644
--- a/suggestions/cogs/help_guild_cog.py
+++ b/suggestions/cogs/help_guild_cog.py
@@ -76,26 +76,16 @@ async def instance_info(
         shard_id = self.bot.get_shard_id(guild_id)
         cluster_id = (
             1
-            if shard_id < 5
-            else 2
             if shard_id < 10
-            else 3
-            if shard_id < 15
-            else 4
+            else 2
             if shard_id < 20
-            else 5
-            if shard_id < 25
-            else 6
+            else 3
             if shard_id < 30
-            else 7
-            if shard_id < 35
-            else 8
+            else 4
             if shard_id < 40
-            else 9
-            if shard_id < 45
-            else 10
+            else 5
             if shard_id < 50
-            else 11
+            else 6
         )
 
         await interaction.send(
@@ -164,7 +154,7 @@ async def show_bot_status(
             title="Bot infrastructure status",
         )
         down_shards: list[str] = [str(i) for i in range(53)]
-        down_clusters: list[str] = [str(i) for i in range(1, 12)]
+        down_clusters: list[str] = [str(i) for i in range(1, 7)]
         avg_bot_latency: list[float] = []
         async with aiohttp.ClientSession(
             headers={"X-API-KEY": os.environ["GARVEN_API_KEY"]}
diff --git a/suggestions/main.py b/suggestions/main.py
index 602c79c..f1d6074 100644
--- a/suggestions/main.py
+++ b/suggestions/main.py
@@ -29,7 +29,7 @@ async def create_bot(database_wrapper=None) -> SuggestionsBot:
         total_shards = 53
         cluster_id = int(os.environ["CLUSTER"])
         offset = cluster_id - 1
-        number_of_shards_per_cluster = 5
+        number_of_shards_per_cluster = 10
         shard_ids = [
             i
             for i in range(
diff --git a/suggestions/zonis_routes.py b/suggestions/zonis_routes.py
index 5e9b284..f2bd73c 100644
--- a/suggestions/zonis_routes.py
+++ b/suggestions/zonis_routes.py
@@ -37,7 +37,6 @@ def __init__(self, bot: SuggestionsBot):
             "share_with_devs",
             "refresh_premium",
             "shared_guilds",
-            "cached_item_count",
         )
 
     async def start(self):
@@ -81,22 +80,10 @@ async def share_with_devs(self, title, description, sender):
         embed.set_footer(text=f"Sender: {sender}")
         await channel.send(embed=embed)
 
-    @client.route()
-    async def cached_item_count(self) -> dict[str, int]:
-        state = self.bot.state
-        stats = self.bot.stats
-        suggestions_queue_cog = self.bot.get_cog("SuggestionsQueueCog")
-        data = {
-            "state.autocomplete_cache": len(state.autocomplete_cache),
-            "state.guild_cache": len(state.guild_cache),
-            "state.view_voters_cache": len(state.view_voters_cache),
-            "state.guild_configs": len(state.guild_configs),
-            "state.user_configs": len(state.user_configs),
-            "stats.cluster_guild_cache": len(stats.cluster_guild_cache),
-            "stats.member_stats_cache": len(stats.member_stats_cache),
-            "suggestions_queue_cog.paginator_objects": len(
-                suggestions_queue_cog.paginator_objects  # noqa
-            ),
-        }
+    async def refresh_premium(self, user_id: int):
+        # TODO Implement the ability to refresh premium states for a user
+        return True
 
-        return data
+    @client.route()
+    async def shared_guilds(self, guild_ids: list[int]):
+        return [gid for gid in guild_ids if gid in self.bot.guild_ids]