1 files changed, 88 insertions, 57 deletions

diff --git a/cogs/compression.py b/cogs/compression.py
index f39b56a..b684f9f 100644
--- a/cogs/compression.py
+++ b/cogs/compression.py
@@ -1,6 +1,7 @@
 import mimetypes
 import os
 import subprocess
+from typing import List, Optional
 
 import discord
 from discord.ext import commands
@@ -19,21 +20,95 @@ def filetype(filename):
             return "Image"
     return "Unknown"
 
-
 class Compression(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
+        self.temp_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "temp")
+        if not os.path.exists(self.temp_dir):
+            os.makedirs(self.temp_dir)
+
+    async def _compress_attachment(self, attachment: discord.Attachment) -> Optional[discord.File]:
+        """Compress a single attachment and return the compressed file"""
+        try:
+            file_path = os.path.join(self.temp_dir, attachment.filename)
+            output_path = os.path.join(self.temp_dir, f"better_{attachment.filename}")
+            
+            # Save the original file
+            await attachment.save(file_path)
+            
+            # Process the file
+            file_type = filetype(file_path)
+            logger.info(f"Compression started for: {file_type}")
+            
+            if file_type == "Video":
+                command = f"ffmpeg -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":
+                command = f"ffmpeg -i {file_path} -ac 1 -b:a 500 -acodec libopus {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:
+                logger.info(f"Unsupported file type: {file_type}")
+                os.remove(file_path)
+                return None
+            
+            # Run the compression
+            subprocess.run(command, check=True, shell=True)
+            
+            if os.path.exists(output_path):
+                logger.info(f"Successfully compressed: {attachment.filename}")
+                return discord.File(output_path, filename=f"compressed_{attachment.filename}")
+            return None
+            
+        except Exception as e:
+            logger.error(f"Error processing {attachment.filename}: {str(e)}")
+            return None
+            
+        finally:
+            # Clean up temporary files
+            if 'file_path' in locals() and os.path.exists(file_path):
+                os.remove(file_path)
+
+    async def _compress_attachments(self, attachments: List[discord.Attachment]) -> List[discord.File]:
+        """Compress multiple attachments and return list of compressed files"""
+        compressed_files = []
+        for attachment in attachments:
+            compressed_file = await self._compress_attachment(attachment)
+            if compressed_file:
+                compressed_files.append(compressed_file)
+        return compressed_files
 
     @commands.message_command(
         name="kompresja",
         description="skompresuj wideło albo dźwienk :)",
     )
     async def message_compression(self, ctx, message: discord.Message):
-        if message.attachments:
-            for attachment in message.attachments:
-                await self.compression(ctx, attachment)
-        else:
+        if not message.attachments:
             await ctx.respond("Brak załączników w wiadomości", ephemeral=True)
+            return
+            
+        await ctx.defer()
+        compressed_files = await self._compress_attachments(message.attachments)
+        
+        # Send all compressed files in one message
+        if compressed_files:
+            await ctx.respond(
+                f"Skompresowano {len(compressed_files)} z {len(message.attachments)} plików:",
+                files=compressed_files
+            )
+            
+            # Clean up compressed files
+            for file in compressed_files:
+                if os.path.exists(file.fp.name):
+                    os.remove(file.fp.name)
+        else:
+            await ctx.respond("Nie udało się skompresować żadnych plików.", ephemeral=True)
 
     @commands.slash_command(
         name="kompresja",
@@ -45,60 +120,16 @@ class Compression(commands.Cog):
         required=True,
         type=discord.Attachment
     )
-    async def compression(self, ctx, plik):
+    async def compression(self, ctx, plik: discord.Attachment):
         await ctx.defer()
-        # Save files to the temp directory inside the bot's root directory
-        bot_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-        temp_dir = os.path.join(bot_dir, "temp")
-        if not os.path.exists(temp_dir):
-            os.makedirs(temp_dir)
-        file_path = os.path.join(temp_dir, plik.filename)
-        await plik.save(file_path)
-
-        input_path = file_path
-        output_path = os.path.join(temp_dir, f"better_{plik.filename}")
-        file_type = filetype(file_path)
-        logger.info(f"Compression started for: {file_type}")
-        if file_type == "Video":
-            # Execute ffmpeg command to compress video
-            ffmpeg_command = f"ffmpeg -i {input_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":
-            ffmpeg_command = f"ffmpeg -i {input_path} -ac 1 -b:a 500 -acodec libopus {output_path}"
-        elif file_type == "Image":
-            # Use ImageMagick's magick to heavily compress image (JPEG/PNG/GIF)
-            # For JPEG: quality 5, strip metadata, resize to 30% of the original size
-            # For PNG: magick to JPEG, then compress
-            # For GIF: optimize and reduce colors
-            ext = os.path.splitext(input_path)[1].lower()
-            if ext in [".jpg", ".jpeg"]:
-                img_command = f"magick {input_path} -strip -resize 30% -quality 5 {output_path}"
-            elif ext == ".gif":
-                # Optimize GIF: reduce colors, optimize layers, and resize
-                img_command = f"magick {input_path} -coalesce -resize 30% -colors 32 -strip -layers Optimize {output_path}"
-            else:
-                # magick non-JPEG images to JPEG for maximum compression
-                output_path = os.path.splitext(output_path)[0] + ".jpg"
-                img_command = f"magick {input_path} -strip -resize 30% -quality 5 {output_path}"
-            ffmpeg_command = img_command
+        compressed_file = await self._compress_attachment(plik)
+        
+        if compressed_file:
+            await ctx.respond(file=compressed_file)
+            if os.path.exists(compressed_file.fp.name):
+                os.remove(compressed_file.fp.name)
         else:
-            await ctx.respond("Plik jest nieobsługiwany", ephemeral=True)
-            logger.info("Unsupported file type")
-            return
-
-        try:
-            subprocess.run(ffmpeg_command, check=True, shell=True)
-            # Upload the compressed file
-            await ctx.respond(file=discord.File(output_path))
-            logger.info(f"Compression completed for: {file_type}")
-            # Delete compressed file
-            os.remove(input_path)
-            os.remove(output_path)
-        except subprocess.CalledProcessError as e:
-            await ctx.respond(f"Compression failed: {str(e)}", ephemeral=True)
-            if os.path.exists(input_path):
-                os.remove(input_path)
-            if os.path.exists(output_path):
-                os.remove(output_path)
+            await ctx.respond("Nie udało się skompresować pliku.", ephemeral=True)
 
 
 def setup(bot):