1 files changed, 55 insertions, 21 deletions

diff --git a/cogs/compression.py b/cogs/compression.py
index d235c6d..1fcefb5 100644
--- a/cogs/compression.py
+++ b/cogs/compression.py
@@ -5,7 +5,7 @@
 
 import mimetypes
 import os
-import subprocess
+import asyncio
 from typing import List, Optional
 from os import PathLike
 import discord
@@ -32,6 +32,38 @@ class Compression(commands.Cog):
         if not os.path.exists(self.temp_dir):
             os.makedirs(self.temp_dir)
 
+    def _process_video(self, file_path: str, output_path: str) -> str:
+        """Process video compression"""
+        command = f"ffmpeg -y -i {file_path} -c:v libx264 -crf 51 -b:v 100 -vf 'fps=10,scale=trunc(iw*0.5/2)*2:trunc(ih*0.5/2)*2' -preset ultrafast -c:a libopus -b:a 500 {output_path}"
+        return command
+    
+    def _process_audio(self, file_path: str, output_path: str) -> str:
+        """Process audio compression"""
+        if output_path.lower().endswith('.mp3'):
+            output_path = os.path.splitext(output_path)[0] + os.path.splitext(output_path)[1].replace('.mp3', '')
+        if not output_path.lower().endswith('.ogg'):
+            output_path = os.path.splitext(output_path)[0] + '.ogg'
+        command = f"ffmpeg -y -i {file_path} -c:a libopus -b:a 500 -ar 12000 -ac 1 {output_path}"
+        return command
+    
+    def _process_image(self, file_path: str, output_path: str) -> tuple[str, str]:
+        """Process image compression"""
+        ext = os.path.splitext(file_path)[1].lower()
+        if ext in [".jpg", ".jpeg"]:
+            command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}"
+        elif ext == ".gif":
+            command = f"magick {file_path} -coalesce -resize 30% -colors 32 -strip -layers Optimize {output_path}"
+        else:
+            output_path = os.path.splitext(output_path)[0] + ".jpg"
+            command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}"
+        return command, output_path
+    
+    async def _execute_compression(self, command: str) -> bool:
+        """Execute compression command"""
+        proc = await asyncio.create_subprocess_exec(command)
+        await proc.wait()
+        return proc.returncode == 0
+    
     async def _compress_attachment(self, attachment: discord.Attachment) -> Optional[discord.File]:
         """Compress a single attachment and return the compressed file"""
         file_path = PathLike[str]
@@ -46,32 +78,23 @@ class Compression(commands.Cog):
             file_type = filetype(file_path)
             logger.info(f"Compression started for: {file_type}")
             
-            if file_type == "Video":
-                command = f"ffmpeg -y -i {file_path} -c:v libx264 -crf 51 -b:v 100 -vf 'fps=10,scale=trunc(iw*0.5/2)*2:trunc(ih*0.5/2)*2' -preset ultrafast -c:a libopus -b:a 500 {output_path}"
-            elif file_type == "Audio":
-                if output_path.lower().endswith('.mp3'):
-                    output_path = os.path.splitext(output_path)[0] + os.path.splitext(output_path)[1].replace('.mp3', '')
-                if not output_path.lower().endswith('.ogg'):
-                    output_path = os.path.splitext(output_path)[0] + '.ogg'
-                command = f"ffmpeg -y -i {file_path} -c:a libopus -b:a 500 -ar 12000 -ac 1 {output_path}"
-            elif file_type == "Image":
-                ext = os.path.splitext(file_path)[1].lower()
-                if ext in [".jpg", ".jpeg"]:
-                    command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}"
-                elif ext == ".gif":
-                    command = f"magick {file_path} -coalesce -resize 30% -colors 32 -strip -layers Optimize {output_path}"
-                else:
-                    output_path = os.path.splitext(output_path)[0] + ".jpg"
-                    command = f"magick {file_path} -strip -resize 30% -quality 5 {output_path}"
-            else:
+            command = self._get_compression_command(file_type, file_path, output_path)
+            if command is None:
                 logger.info(f"Unsupported file type: {file_type}")
                 os.remove(file_path)
                 return None
             
+            # Update output_path if it was modified
+            if file_type == "Image":
+                command, output_path = command
+            elif file_type == "Audio":
+                # output_path is already modified in _process_audio
+                pass
+            
             # Run the compression
-            subprocess.run(command, check=True, shell=True)
+            success = await self._execute_compression(command)
             
-            if os.path.exists(output_path):
+            if success and os.path.exists(output_path):
                 logger.info(f"Successfully compressed: {attachment.filename}")
                 return discord.File(output_path, filename=f"compressed_{attachment.filename}")
             return None
@@ -84,6 +107,17 @@ class Compression(commands.Cog):
             # Clean up temporary files
             if 'file_path' in locals() and os.path.exists(file_path):
                 os.remove(file_path)
+    
+    def _get_compression_command(self, file_type: str, file_path: str, output_path: str):
+        """Get the appropriate compression command based on file type"""
+        if file_type == "Video":
+            return self._process_video(file_path, output_path)
+        elif file_type == "Audio":
+            return self._process_audio(file_path, output_path)
+        elif file_type == "Image":
+            return self._process_image(file_path, output_path)
+        else:
+            return None
 
     async def _compress_attachments(self, attachments: List[discord.Attachment]) -> List[discord.File]:
         """Compress multiple attachments and return list of compressed files"""