path: root/cogs/compression.py

# This file is part of Senator Sperma
# Licensed under the GPL-3.0 License
# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
# https://codeberg.org/salmonik/senator-sperma

import mimetypes
import os
import asyncio
import aiofiles
import requests
import urllib.parse
from typing import List, Optional
from os import PathLike
import discord
from discord.ext import commands
from discord import AutocompleteContext
from cogs.utils import file_search as search_utils
import shlex
import logging

logger = logging.getLogger(__name__)

def filetype(filename):
    mime_type, _ = mimetypes.guess_type(filename)
    if mime_type:
        if mime_type.startswith('video/'):
            return "Video"
        elif mime_type.startswith('audio/'):
            return "Audio"
        elif mime_type.startswith('image/'):
            return "Image"
    return "Unknown"

class Compression(commands.Cog):
    def __init__(self, bot):
        from cogs.utils.temp import get_temp_dir
        from cogs.utils.settings import get_baseurl
        self.bot = bot
        self.temp_dir = get_temp_dir()
        self._cached_urls = {}
        self._baseurl = get_baseurl()

    async def _file_autocomplete(self, ctx: AutocompleteContext):
        query = ctx.value.lower() if ctx.value else ""
        results = await search_utils.UniversalSearcher.search_copyparty_autocomplete(query, max_results=25)

        user_id = ctx.interaction.user.id
        self._cached_urls[user_id] = {urllib.parse.unquote(url.split('/')[-1]): url for url in results}

        choices = []
        for url in results[:25]:
            filename = urllib.parse.unquote(url.split('/')[-1])
            display_name = filename[:99] + "…" if len(filename) > 100 else filename
            choices.append(discord.OptionChoice(name=display_name, value=filename))
        return choices

    async def _download_from_url(self, url: str) -> Optional[str]:
        """Download file from copyparty URL and return local path"""
        try:
            response = requests.get(url)
            if response.status_code != 200:
                return None

            filename = url.split('/')[-1].split('?')[0]
            file_path = os.path.join(self.temp_dir, filename)

            async with aiofiles.open(file_path, 'wb') as f:
                await f.write(response.content)
            return file_path
        except Exception as e:
            logger.error(f"Error downloading from {url}: {e}")
            return None

    async def _upload_to_copyparty(self, file_path: str) -> Optional[str]:
        """Upload file to copyparty and return download URL"""
        try:
            async with aiofiles.open(file_path, 'rb') as f:
                file_content = await f.read()

            files = {'f': (os.path.basename(file_path), file_content, 'application/octet-stream')}
            headers = {'Cookie': f'cppwd={os.getenv("COPYPARTY_PASSWORD")}'}
            params = {'j': '1', 'ck': 'b2'}

            response = requests.post(self._baseurl, files=files, headers=headers, params=params)

            if response.status_code in [200, 201]:
                result = response.json()
                return result.get('fileurl')
        except Exception as e:
            logger.error(f"Upload error: {e}")
            return None

    def _process_video(self, file_path: str, output_path: str) -> str:
        """Process video compression"""
        command = f"ffmpeg -y -i {shlex.quote(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 {shlex.quote(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 {shlex.quote(file_path)} -c:a libopus -b:a 500 -ar 12000 -ac 1 {shlex.quote(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 {shlex.quote(file_path)} -strip -resize 30% -quality 5 {shlex.quote(output_path)}"
        elif ext == ".gif":
            command = f"magick {shlex.quote(file_path)} -coalesce -resize 30% -colors 32 -strip -layers Optimize {shlex.quote(output_path)}"
        else:
            output_path = os.path.splitext(output_path)[0] + ".jpg"
            command = f"magick {shlex.quote(file_path)} -strip -resize 30% -quality 5 {shlex.quote(output_path)}"
        return command, output_path

    async def _execute_compression(self, command: str) -> bool:
        """Execute compression command"""
        proc = await asyncio.create_subprocess_exec(
            *shlex.split(command),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        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"""
        try:
            file_path: str = 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}")

            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
            success = await self._execute_compression(command)

            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

        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)

    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"""
        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 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):  # type: ignore[attr-defined]
                    os.remove(file.fp.name)  # type: ignore[attr-defined]
        else:
            await ctx.respond("Nie udało się skompresować żadnych plików.", ephemeral=True)

    @commands.slash_command(
        name="kompresja",
        description="skompresuj wideło albo dźwienk :)",
    )
    @discord.option(
        name="plik",
        description="plik z copyparty",
        required=False,
        autocomplete=_file_autocomplete
    )
    @discord.option(
        name="zalacznik",
        description="plik z discorda",
        required=False,
        type=discord.Attachment
    )
    @discord.option(
        name="serwer?",
        description="wysłać plik na serwer?",
        required=False,
        default=False,
        type=bool
    )
    async def compression(self, ctx, plik: str = None, zalacznik: discord.Attachment = None, server = False):
        await ctx.defer()

        file_path = None
        original_filename = None

        if plik:
            url = self._cached_urls.get(ctx.author.id, {}).get(plik)
            if not url:
                await ctx.respond("Wybierz plik z podpowiedzi", ephemeral=True)
                return
            file_path = await self._download_from_url(url)
            original_filename = plik
        elif zalacznik:
            file_path = os.path.join(self.temp_dir, zalacznik.filename)
            original_filename = zalacznik.filename
            await zalacznik.save(file_path)
        else:
            await ctx.respond("Podaj plik z copyparty lub załącznik", ephemeral=True)
            return

        if not file_path or not os.path.exists(file_path):
            await ctx.respond("Nie udało się pobrać pliku", ephemeral=True)
            return

        try:
            file_type = filetype(file_path)
            output_path = os.path.join(self.temp_dir, f"better_{original_filename}")

            command = self._get_compression_command(file_type, file_path, output_path)
            if command is None:
                await ctx.respond("Nieobsługiwany typ pliku", ephemeral=True)
                return

            if file_type == "Image":
                command, output_path = command
            elif file_type == "Audio":
                # Audio output path remains unchanged - already processed in _process_audio
                pass

            success = await self._execute_compression(command)

            if not success or not os.path.exists(output_path):
                await ctx.respond("Kompresja nie powiodła się", ephemeral=True)
                return

            file_size = os.path.getsize(output_path)
            if file_size > 1024 * 1024 * 10 or server:
                dl_url = await self._upload_to_copyparty(output_path)
                if dl_url:
                    await ctx.respond(f"Skompresowano\n{dl_url}")
                else:
                    await ctx.respond("Nie udało wysłać na serwer", ephemeral=True)
            else:
                await ctx.respond(file=discord.File(output_path, filename=f"compressed_{original_filename}"))

        finally:
            if file_path and os.path.exists(file_path):
                os.remove(file_path)
            if 'output_path' in locals() and os.path.exists(output_path):
                os.remove(output_path)


def setup(bot):
    bot.add_cog(Compression(bot))