Refactor UsrBg classes and improve caching mechanism for user backgro…

…unds

classes/usrbg.py

@@ -1,112 +1,77 @@
-import json
 from dataclasses import dataclass
+from typing import Any
 
 import aiohttp
-from fake_useragent import FakeUserAgent
 from interactions import Snowflake
 
 from classes.cache import Caching
+from modules.const import USER_AGENT
 
-USER_AGENT = FakeUserAgent(browsers=["chrome", "edge", "opera"]).random
-Cache = Caching(cache_directory="cache/usrbg", cache_expiration_time=216000)
+# 2 days and half
+day2h = 60 * 60 * 60
+Cache = Caching(cache_directory="cache/usrbg", cache_expiration_time=day2h)
+
+BASE_URL = "https://usrbg.is-hardly.online"
 
 
 @dataclass
-class UserBackgroundStruct:
-    """A dataclass to represent user background."""
+class UsrBgDataStruct:
+    """A dataclass pointing to the user's background"""
+
+    users: dict[str, str]
+    endpoint: str = BASE_URL
+    bucket: str = "usrbg"
+    prefix: str = "v2/"
+
+    def get_user_banner(self, user_id: Snowflake) -> str:
+        """Get the user's background from the cache or the API"""
 
-    _id: str
-    """Entry ID"""
-    uid: Snowflake
-    """User ID"""
-    img: str
-    """Image URL"""
-    orientation: str
-    """Image orientation"""
+        if str(user_id) in self.users:
+            loc = self.users[str(user_id)]
+            return f"{self.endpoint}/{self.bucket}/{self.prefix}{str(user_id)}?{loc}"
+
+        raise ValueError(
+            "User not found in the UsrBG Database. Either is cache is old, or the user has no background."
+        )
 
 
-class UserBackground:
-    """usrbg wrapper"""
+class UsrBg:
+    """A class to interact with the UsrBG API"""
 
     def __init__(self):
-        """Initialize the UserBackground class."""
-        self.raw_url = "https://raw.githubusercontent.com/Discord-Custom-Covers/usrbg/master/dist/usrbg.json"
+        """Initialize the class"""
+        self.headers = {"User-Agent": USER_AGENT}
         self.session = None
-        self.headers = None
+        self.database: dict[str, str] = {}
 
     async def __aenter__(self):
-        """Enter the async context manager."""
-        self.headers = {"User-Agent": USER_AGENT}
+        """Return the class instance"""
         self.session = aiohttp.ClientSession(headers=self.headers)
         return self
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        """Exit the async context manager."""
-        await self.close()
-
-    async def close(self):
-        """Close the aiohttp session."""
-        await self.session.close()
-
-    async def _fetch_background(self) -> dict:
-        """
-        Fetch the background from the GitHub repository.
-
-        Returns:
-            dict: The background data.
-        """
-        print("Fetching usrbg from GitHub...")
-        async with self.session.get(self.raw_url) as response:
-            resp = await response.text()
-            data = json.loads(resp)
-            return data
-
-    @staticmethod
-    async def _find_user(user_id: Snowflake, data: dict) -> UserBackgroundStruct | None:
-        """
-        Find user on the dict, then return datastruct
-
-        Args:
-            user_id (Snowflake): User's Discord ID
-            data: usrbg's dict
-
-        Return:
-            UserBackgroundStruct: Datastruct
-            None: If user can't be found
-        """
-        # Find user
-        user = next(
-            (item for item in data if item["uid"] == str(user_id)),
-            None)
-        if user is None:
-            return None
-        # Create datastruct
-        return UserBackgroundStruct(
-            _id=user["_id"],
-            uid=user["uid"],
-            img=user["img"],
-            orientation=user["orientation"],
-        )
-
-    async def get_background(self, user_id: Snowflake) -> UserBackgroundStruct | None:
-        """
-        Get the user background.
-
-        Args:
-            user_id (Snowflake): User's Discord ID
+    async def __aexit__(self, exc_type, exc_value, traceback):  # type: ignore
+        """Close the session"""
+        await self.session.close() if self.session else None
 
-        Returns:
-            UserBackgroundStruct: The user background.
-            None: If user can't be found
-        """
-        # Get cache path
+    async def _fetch_background(self) -> Any:
+        """Get the user's background from the API"""
         cache_path = Cache.get_cache_path("usrbg.json")
-        # Read cache
-        data = Cache.read_cache(cache_path)
-        # If cache is expired, fetch from GitHub
-        if data is None:
-            data = await self._fetch_background()
-            Cache.write_cache(cache_path, data)
-        # Find user
-        user = await self._find_user(user_id, data)
-        return user
+        self.database = Cache.read_cache(cache_path)
+        if self.database:
+            return self.database
+
+        async with aiohttp.ClientSession(headers=self.headers) as session:
+            async with session.get(f"{BASE_URL}/users") as resp:
+                if resp.status == 200:
+                    self.database = await resp.json()
+                    Cache.write_cache(cache_path, self.database)
+                    return self.database
+                resp.raise_for_status()
+
+    async def get_background(self, user_id: Snowflake) -> str:
+        """Get the user's background from the API"""
+        if not self.database:
+            await self._fetch_background()
+
+        resp = UsrBgDataStruct(**self.database).get_user_banner(user_id)
+        return resp