Overhaul + Sonarr support!

5 files changed, 364 insertions, 129 deletions

diff --git a/code/cogs/error.py b/code/cogs/error.py
index 2752400..6b6e557 100644
--- a/code/cogs/error.py
+++ b/code/cogs/error.py
@@ -15,10 +15,16 @@ class slash_handlers(commands.Cog):
         ):
             embed = discord.Embed(
                 title="Jellyfin Account Creation Disabled",
-                description=f"The owner of {self.bot.user.mention} has disabled the ability to create temporary Jellyfin accounts. Contact an administrator for more information.",
-                color=0xD01B86
+                description=(
+                    f"The owner of {self.bot.user.mention} has disabled the"
+                    " ability to create temporary Jellyfin accounts. Contact"
+                    " an administrator for more information."
+                ),
+                color=0xD01B86,
+            )
+            await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            await interaction.response.send_message(embed=embed, ephemeral=True)
         else:
             raise error
 
diff --git a/code/cogs/newaccount.py b/code/cogs/newaccount.py
index db5abd0..b341147 100644
--- a/code/cogs/newaccount.py
+++ b/code/cogs/newaccount.py
@@ -3,8 +3,12 @@ from discord import app_commands
 from discord.ext import commands
 import sqlite3
 
-from func.jellyfin import create_jellyfin_account
-from global_variables import JELLYFIN_URL, ENABLE_JELLYFIN_TEMP_ACCOUNTS, ACCOUNT_TIME
+from utils.jellyfin_create import create_jellyfin_account
+from utils.config import (
+    JELLYFIN_URL,
+    JELLYFIN_ENABLED,
+    ACCOUNT_TIME,
+)
 
 
 class NewAccount(commands.Cog):
@@ -12,47 +16,77 @@ class NewAccount(commands.Cog):
         self.bot = bot
 
     @app_commands.command()
-    @app_commands.check(lambda inter: ENABLE_JELLYFIN_TEMP_ACCOUNTS)
-    async def newaccount(self, interaction: discord.Interaction):
-        "Create a new temporary Jellyfin account"
+    @app_commands.check(lambda inter: JELLYFIN_ENABLED)
+    async def newaccount(self, interaction: discord.Interaction) -> None:
+        """Create a new temporary Jellyfin account"""
         # Make sure the user doesn't already have an account
         db = sqlite3.connect("cordarr.db")
         cursor = db.cursor()
         cursor.execute(
-            "SELECT * FROM jellyfin_accounts WHERE user_id = ?", (interaction.user.id,)
+            "SELECT * FROM jellyfin_accounts WHERE user_id = ?",
+            (interaction.user.id,),
         )
-        if cursor.fetchone():
+        account = cursor.fetchone()
+        db.close()
+        # Account already allocated
+        if account:
             embed = discord.Embed(
                 title="Account Already Exists",
-                description="Look at your previous DMs with me to find your account information. You will be permitted to create a new account after your current one expires.",
-                color=0xD01B86
+                description=(
+                    "Look at your previous DMs with me to find your account"
+                    " information. You will be permitted to create a new"
+                    " account after your current one expires."
+                ),
+                color=0xD01B86,
+            )
+            return await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
 
         # Create a new Jellyfin account for the user
         response = create_jellyfin_account(interaction.user.id)
         if response:
             embed = discord.Embed(
                 title="Account Created",
-                description="Your account has been successfully created. Check your DMs for your account information.",
-                color=0xD01B86
+                description=(
+                    "Your account has been successfully created. Check your"
+                    " DMs for your account information."
+                ),
+                color=0xD01B86,
+            )
+            await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            await interaction.response.send_message(embed=embed, ephemeral=True)
 
             # Send the user their account information
             embed = discord.Embed(
                 title="Jellyfin Account Information",
-                description=f"Here is your temporary account information. You will need this to access the Jellyfin server.\n\n**Server URL:** `{JELLYFIN_URL}`\n**Username:** `{response[0]}`\n**Password:** `{response[1]}`\n\nYour account will be automatically deleted in {ACCOUNT_TIME} hours.",
-                color=0xD01B86
+                description=(
+                    # fmt: off
+                    "Here is your temporary account information.\n\n"
+                    f"**Server URL:** `{JELLYFIN_URL}`\n"
+                    f"**Username:** `{response[0]}`\n"
+                    f"**Password:** `{response[1]}`\n\n"
+                    "Your account will be automatically deleted in"
+                    f" {ACCOUNT_TIME} hours."
+                    # fmt: on
+                ),
+                color=0xD01B86,
             )
             await interaction.user.send(embed=embed)
+        # If account not created for some reason
         else:
             embed = discord.Embed(
                 title="Unknown Error Occured",
-                description="Error creating Jellyfin account. Please try again. If the error persists, contact an administrator.",
-                color=0xD01B86
+                description=(
+                    "Error creating Jellyfin account. Please try again. If the"
+                    " error persists, contact an administrator."
+                ),
+                color=0xD01B86,
+            )
+            return await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
 
 
 async def setup(bot):
diff --git a/code/cogs/request.py b/code/cogs/request.py
index fb1706c..2ad71ba 100644
--- a/code/cogs/request.py
+++ b/code/cogs/request.py
@@ -1,53 +1,111 @@
 import discord
 from discord import app_commands
 from discord.ext import commands
+from typing import Literal
 
-from func.radarr import get_movies, AddMovieView
+from utils.content_get import get_content
+from utils.content_view import AddContentView
+from utils.config import (
+    RADARR_HOST_URL,
+    RADARR_HEADERS,
+    RADARR_ROOT_FOLDER_PATH,
+    RADARR_QUALITY_PROFILE_ID,
+    SONARR_HOST_URL,
+    SONARR_HEADERS,
+    SONARR_ROOT_FOLDER_PATH,
+    SONARR_QUALITY_PROFILE_ID,
+)
 
 
-class Request(commands.GroupCog, name="request"):
+class Request(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
 
-    @app_commands.command(name="movie")
-    @app_commands.describe(name="Name of the movie to add")
-    async def request_movie(self, interaction: discord.Interaction, name: str):
-        "Request a movie to be added to the Radarr library"
-        movie_data = get_movies(name)
-        if movie_data == "NO RESULTS":
+    @app_commands.command()
+    @app_commands.describe(form="Are you requesting a Movie or Show?")
+    @app_commands.describe(name="Name of the content")
+    async def request(
+        self,
+        interaction: discord.Interaction,
+        form: Literal["Movie", "Show"],
+        name: str,
+    ) -> None:
+        """Request a movie or tv show to be added to the library"""
+        # Get matching content from relevant service
+        if form == "Movie":
+            content_data = get_content(
+                name, "radarr", RADARR_HOST_URL, RADARR_HEADERS
+            )
+        else:
+            content_data = get_content(
+                name, "sonarr", SONARR_HOST_URL, SONARR_HEADERS
+            )
+
+        if content_data == "NO RESULTS":
             embed = discord.Embed(
                 title="No Results",
-                description="No results were found for the given movie name. If you are unable to find the movie, contact an administrator to have it added manually.",
-                color=0xD01B86
+                description=(
+                    # fmt: off
+                    "No results found, please try again. Here are some tips:\n\n"
+                    "1. Double check spelling\n"
+                    "2. Add release year to the query\n"
+                    "3. Double check the \"Movie\" or \"Show\" option"
+                    # fmt: on
+                ),
+                color=0xD01B86,
+            )
+            return await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
 
-        if movie_data == "ALREADY ADDED":
+        if content_data == "ALREADY ADDED":
             embed = discord.Embed(
                 title="Already Added",
-                description="The movie you are trying to add has already been added to the Radarr library.\n\nYou can check the download status of your requests movies by running the `/status` command.",
-                color=0xD01B86
+                description=(
+                    f"**{name}** is already added to the"
+                    f" {'radarr' if form == 'Movie' else 'sonarr'} library. It"
+                    " may be downloading, stalled, or not found. Check the"
+                    " status of the content you have requested with"
+                    " `/status`."
+                ),
+                color=0xD01B86,
+            )
+            return await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
 
         embed = discord.Embed(
             title="Results Found",
-            description="Please select the movie you would like to add from the dropdown below.",
-            color=0xD01B86
+            description=(
+                f"Please select from the top {len(content_data)} results from"
+                f" {'radarr' if form == 'Movie' else 'sonarr'} in the"
+                " dropdown below."
+            ),
+            color=0xD01B86,
         )
-        view = AddMovieView(movie_data)
-        await interaction.response.send_message(embed=embed, view=view, ephemeral=True)
+        # Create view with the content data and relevant service info
+        if form == "Movie":
+            view = AddContentView(
+                content_data,
+                "radarr",
+                RADARR_HOST_URL,
+                RADARR_HEADERS,
+                RADARR_ROOT_FOLDER_PATH,
+                RADARR_QUALITY_PROFILE_ID,
+            )
+        else:
+            view = AddContentView(
+                content_data,
+                "sonarr",
+                SONARR_HOST_URL,
+                SONARR_HEADERS,
+                SONARR_ROOT_FOLDER_PATH,
+                SONARR_QUALITY_PROFILE_ID,
+            )
 
-    @app_commands.command(name="show")
-    @app_commands.describe(name="Name of the show/series to add")
-    async def request_show(self, interaction: discord.Interaction, name: str):
-        "Request a show/series to be added to the Sonarr library"
-        embed = discord.Embed(
-            title="Coming Soon",
-            description="This feature is not yet implemented. Check back later.",
-            color=0xD01B86
+        await interaction.response.send_message(
+            embed=embed, view=view, ephemeral=True
         )
-        await interaction.response.send_message(embed=embed, ephemeral=True)
 
 
 async def setup(bot):
diff --git a/code/cogs/status.py b/code/cogs/status.py
index 7b6b463..abda84a 100644
--- a/code/cogs/status.py
+++ b/code/cogs/status.py
@@ -3,10 +3,13 @@ from discord import app_commands
 from discord.ext import commands
 import requests
 import sqlite3
-import datetime
-import humanize
 
-from global_variables import RADARR_HOST_URL, RADARR_HEADERS
+from utils.config import (
+    RADARR_HOST_URL,
+    RADARR_HEADERS,
+    SONARR_HOST_URL,
+    SONARR_HEADERS,
+)
 
 
 class Status(commands.Cog):
@@ -14,97 +17,229 @@ class Status(commands.Cog):
         self.bot = bot
 
     @app_commands.command()
-    async def status(self, interaction: discord.Interaction):
-        "Get the status of the movies you have requested"
-        # Get all the movie_ids that were requested by the user
+    async def status(self, interaction: discord.Interaction) -> None:
+        """Get the status of the movies you have requested"""
         db = sqlite3.connect("cordarr.db")
         cursor = db.cursor()
         cursor.execute(
-            "SELECT movie_id, movie_title FROM movies WHERE user_id = ?",
+            "SELECT title, release_year, local_id, tmdbid, tvdbid FROM"
+            " requests WHERE user_id = ?",
             (interaction.user.id,),
         )
-        requested_movies = cursor.fetchall()
+        requested_content = cursor.fetchall()
+        db.close()
 
-        users_movies = {}  # Dictionary to store the movies that the user has requested
-        for movie_id, movie_title in requested_movies:
-            users_movies[movie_id] = movie_title
-        # If theres no movies, return a message saying so
-        if not users_movies:
+        # No content requested
+        if len(requested_content) == 0:
             embed = discord.Embed(
-                title="No Movies Requested",
-                description="You have no movies being downloaded at the moment. If you previously added a movie, it is likely that it has finished downloading. If you believe this is an error, please contact an administrator.",
-                color=0xD01B86
+                title="No Content Requested",
+                description=(
+                    "If you believe this is in error, the content you have"
+                    " requested is likely already downloaded."
+                ),
+                color=0xD01B86,
+            )
+            return await interaction.response.send_message(
+                embed=embed, ephemeral=True
             )
-            return await interaction.response.send_message(embed=embed, ephemeral=True)
-        # Otherwise, create the default embed to display the movies being downloaded
+
+        # Create template embed
         embed = discord.Embed(
-            title="Movies Requested",
-            description="Here are the movies you have requested that are currently being downloaded:\n",
-            color=0xD01B86
+            title="Requested Content",
+            description=(
+                "Below are the movies/shows you have requested that are"
+                " currently being downloaded:\n"
+            ),
+            color=0xD01B86,
+        )
+
+        # Unpack the content
+        radarr_content_info, sonarr_content_info = self.unpack_content(
+            requested_content
         )
+        # Get the descriptions and local IDs found in queue
+        radarr_desc, radarr_added_ids = self.process_queue(
+            radarr_content_info, "radarr"
+        )
+        sonarr_desc, sonarr_added_ids = self.process_queue(
+            sonarr_content_info, "sonarr"
+        )
+
+        added_ids = radarr_added_ids + sonarr_added_ids
+        # Get the description of content not in the queue
+        non_queue_desc = self.get_non_queue_content(
+            requested_content, added_ids, interaction.user.id
+        )
+
+        embed.description += radarr_desc + sonarr_desc + non_queue_desc
+
+        await interaction.response.send_message(embed=embed, ephemeral=True)
+
+    def unpack_content(self, requested_content: list) -> tuple:
+        """
+        Given a list of requested content, unpack it into two dictionaries
+
+        Args:
+            requested_content (list): A list of requested content
+
+        Returns:
+            tuple: A tuple of two dictionaries
+        """
+
+        radarr_content_info = {}
+        sonarr_content_info = {}
+
+        for content in requested_content:
+            title, release_year, local_id, tmdbid, tvdbid = content
+            if tmdbid is not None:
+                radarr_content_info[local_id] = {
+                    "title": title,
+                    "release_year": release_year,
+                    "tmdbid": tmdbid,
+                }
+            else:
+                sonarr_content_info[local_id] = {
+                    "title": title,
+                    "release_year": release_year,
+                    "tvdbid": tvdbid,
+                }
+
+        return radarr_content_info, sonarr_content_info
+
+    def process_queue(self, content_info: dict, service: str) -> str:
+        """
+        Given a dictionary of requested content and "sonarr"/"radarr", process the queue
+
+        Args:
+            content_info (dict): A dictionary of content information
+            service (str): The service to check the queue of
 
-        # Now, we get the download status of all movies from the Radarr queue
-        response = requests.get(
-            f"{RADARR_HOST_URL}/api/v3/queue/", headers=RADARR_HEADERS
+        Returns:
+            str: The description of the embed
+        """
+
+        description = ""
+        added_ids = []
+
+        queue = requests.get(
+            f"{RADARR_HOST_URL if service == 'radarr' else SONARR_HOST_URL}/api/v3/queue",
+            headers=RADARR_HEADERS if service == "radarr" else SONARR_HEADERS,
         ).json()
 
-        count = 0
-        added_movie_ids = []
-        for movie in response["records"]:
-            movie_id = movie["movieId"]
-            # If the movie is user requested and is being downloaded
-            if movie_id in users_movies.keys():
-                count += 1
-                added_movie_ids.append(movie_id)
-                if movie["status"] == "downloading":
-                    # Humanize the download time left, or result to 'Unknown
-                    try:
-                        time_left = humanize.precisedelta(
-                            datetime.datetime.strptime(movie["timeleft"], "%H:%M:%S")
-                            - datetime.datetime.strptime("00:00:00", "%H:%M:%S"),
-                            minimum_unit="seconds",
-                        )
-                    except ValueError:
-                        # Sometimes movies will download extremely show and therefore might
-                        # show 'days' in the time left, so strptime appropriately
-                        time_left = humanize.precisedelta(
-                            datetime.datetime.strptime(movie["timeleft"], "%d.%H:%M:%S")
-                            - datetime.datetime.strptime("00:00:00", "%H:%M:%S"),
-                            minimum_unit="seconds",
-                        )
-                    except KeyError or ValueError:
-                        time_left = "Unknown"
-
-                    # Add all the information
-                    embed.description += f"\n{count}. **{users_movies[movie_id]}** - Time Left: ` {time_left} `"
-                else:
-                    embed.description += f"\n{count}. **{users_movies[movie_id]}** - Status: `{str(movie['status']).upper()}`"
-
-        # If a movie wasn't found in the Radarr queue, then it has either finished downloading
-        # or the movie was never found for download
-        if len(added_movie_ids) != len(users_movies.keys()):
-            # Grab all of the "missing" movies to see if a movie is missing or finished downloading
-            response = requests.get(
-                f"{RADARR_HOST_URL}/api/v3/wanted/missing", headers=RADARR_HEADERS
-            ).json()
-            for movie in response["records"]:
-                movie_id = movie["id"]
-                if movie_id in users_movies.keys() and movie_id not in added_movie_ids:
-                    count += 1
-                    added_movie_ids.append(movie_id)
-                    embed.description += f"\n{count}. **{users_movies[movie_id]}** - Status: ` NOT FOUND `"
-        # If there are still movies that haven't been added to the embed, then they
-        # have finished downloading and can be removed from the database
-        for movie_id in users_movies.keys():
-            if movie_id not in added_movie_ids:
-                cursor.execute(
-                    "DELETE FROM movies WHERE user_id = ? AND movie_id = ?",
-                    (interaction.user.id, movie_id),
+        for download in queue["records"]:
+            id_str = "movieId" if service == "radarr" else "seriesId"
+            # If the content was requested by the user
+            if (
+                download[id_str] in content_info.keys()
+                and download[id_str] not in added_ids
+            ):
+                # Append local ID
+                added_ids.append(download[id_str])
+                # Add the download to the embed
+                try:
+                    time_left = self.process_time(download["timeleft"])
+                except KeyError:
+                    time_left = "Unknown"
+                description += (
+                    f"\n**{content_info[download[id_str]]['title']} ({content_info[download[id_str]]['release_year']})**"
+                    f" - Time Left: `{time_left}`"
                 )
-        db.commit()
-        db.close()
 
-        await interaction.response.send_message(embed=embed, ephemeral=True)
+        return description, added_ids
+
+    def get_non_queue_content(
+        self, requested_content: list, added_ids: list, user_id: int
+    ) -> str:
+        """
+        Given a list of requested content and a list of added IDs, return a description of content not in the queue
+
+        Args:
+            requested_content (list): A list of requested content
+            added_ids (list): A list of IDs that are in the queue
+            user_id (int): The ID of the user
+
+        Returns:
+            str: A description of content not in the queue
+        """
+
+        description = ""
+        # For evry piece of content not in the queue, check if it has a file
+        for content in requested_content:
+            title, release_year, local_id, tmdbid, _ = content
+            # If not in queue
+            if local_id not in added_ids:
+                # Pull the movie data from the service
+                if tmdbid is not None:
+                    data = requests.get(
+                        f"{RADARR_HOST_URL}/api/v3/movie/{local_id}",
+                        headers=RADARR_HEADERS,
+                    ).json()
+                else:
+                    data = requests.get(
+                        f"{SONARR_HOST_URL}/api/v3/series/{local_id}",
+                        headers=SONARR_HEADERS,
+                    ).json()
+
+                # If the movie has a file, then it has finished downloading
+                if data.get("hasFile", True):
+                    # Remove from database
+                    db = sqlite3.connect("cordarr.db")
+                    cursor = db.cursor()
+                    cursor.execute(
+                        "DELETE FROM requests WHERE user_id = ? AND"
+                        " local_id = ?",
+                        (user_id, local_id),
+                    )
+                    db.commit()
+                    db.close()
+                # If series and only a portion of episodes have been downloaded
+                if data.get("statistics").get("percentOfEpisodes"):
+                    description += (
+                        f"\n**{title} ({release_year})** - Status: `NOT"
+                        " FOUND"
+                        f" ({int(data['statistics']['percentOfEpisodes'])}%"
+                        " of eps.)`"
+                    )
+                # All other scenarios, download not found
+                else:
+                    description += (
+                        f"\n**{title} ({release_year})** - Status: `NOT FOUND`"
+                    )
+
+        return description
+
+    def process_time(self, time) -> str:
+        """
+        Given a time string, process it into a human readable format
+
+        Args:
+            time (str): A string representing time
+
+        Returns:
+            str: A human readable time
+        """
+        # Split the input by either ':' or spaces
+        parts = time.replace(" ", ":").replace(".", ":").split(":")
+
+        # Handle different input lengths
+        if len(parts) == 2:  # Format: MM:SS
+            minutes, seconds = map(int, parts)
+            return f"{minutes} min. {seconds} sec."
+
+        elif len(parts) == 3:  # Format: HH:MM:SS
+            hours, minutes, seconds = map(int, parts)
+            if hours == 0:
+                return f"{minutes} min. {seconds} sec."
+            return f"{hours} hr. {minutes} min."
+
+        elif len(parts) == 4:  # Format: D:HH:MM:SS
+            days, hours, minutes, seconds = map(int, parts)
+            if days == 0:
+                return f"{hours} hr. {minutes} min."
+            return f"{days} days {hours} hr."
+
+        else:
+            return "Unknown"
 
 
 async def setup(bot):
diff --git a/code/cogs/tree_sync.py b/code/cogs/tree_sync.py
index 5050730..b84766e 100644
--- a/code/cogs/tree_sync.py
+++ b/code/cogs/tree_sync.py
@@ -9,7 +9,9 @@ class TreeSync(commands.Cog):
     @commands.command()
     @commands.dm_only()
     @commands.is_owner()
-    async def sync(self, ctx: commands.Context, *, guild: Object = None) -> None:
+    async def sync(
+        self, ctx: commands.Context, *, guild: Object = None
+    ) -> None:
         if not guild or guild == None:
             await self.bot.tree.sync()
             await ctx.author.send("Synced commands globally")