Lynxie/lynxie/commands/image.py

import os
import datetime
import requests
from io import BytesIO

from PIL import Image

import discord
from discord.ext import commands

from lynxie.config import IMAGE_EXTENSIONS, IMAGE_OVERLAYS, ASSETS_PATH
from lynxie.utils import error_message


class Img(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self._overlays = {
            "bubble": Image.open(os.path.join(ASSETS_PATH, "bubble.png")),
            "gang": Image.open(os.path.join(ASSETS_PATH, "gang.png")),
        }

    @commands.command()
    async def overlay(self, ctx, style: str = None):
        start_time = datetime.datetime.now()
        style = style.lower().strip() if style else None
        image_attachments = None

        if ctx.message.attachments:
            image_attachments = ctx.message.attachments[0]
        elif ctx.message.reference and ctx.message.reference.resolved.attachments:
            image_attachments = ctx.message.reference.resolved.attachments[0]
        elif ctx.message.embeds and ctx.message.embeds[0].image:
            image_attachments = ctx.message.embeds[0].image
        else:
            async for message in ctx.guild.get_channel(ctx.channel.id).history(limit=10):
                if message.attachments:
                    image_attachments = message.attachments[0]
                    break
                elif message.embeds and message.embeds[0].image:
                    image_attachments = message.embeds[0].image
                    break

        # Check if image should be processed
        async with ctx.typing():
            if not style or style not in IMAGE_OVERLAYS:
                error = (
                    "That is not a valid option! Valid options are:\n"
                    f"`{', '.join(IMAGE_OVERLAYS)}`"
                )
                await ctx.reply(embed=error_message(error))
                return

            if not image_attachments:
                error = "You need to attach an image to use this command!"
                await ctx.reply(embed=error_message(error))
                return

            # Extracts file extension from filename or url
            if (
                image_attachments.filename
                and not image_attachments.filename.split(".")[-1].lower()
                in IMAGE_EXTENSIONS
            ):
                error = (
                    "Unsupported file type! Supported file types are:\n"
                    f"`{', '.join(IMAGE_EXTENSIONS)}`"
                )
                await ctx.reply(embed=error_message(error))
                return
            elif (
                image_attachments.url
                and not image_attachments.url.split(".")[-1].lower()
                in IMAGE_EXTENSIONS
            ):
                error = (
                    "Unsupported file type! Supported file types are:\n"
                    f"`{', '.join(IMAGE_EXTENSIONS)}`"
                )
                await ctx.reply(embed=error_message(error))
                return

            if (
                image_attachments.size
                and image_attachments.size > 8 * 1024 * 1024
            ):
                error = (
                    "That image is too big! Please use an image that is less than 8MB."
                )
                await ctx.reply(embed=error_message(error))
                return

            if (
                not 0 < image_attachments.width <= 3500
                or not 0 < image_attachments.height <= 3500
            ):
                error = "Image must be at least 1x1 and under 3500x3500!"
                await ctx.reply(embed=error_message(error))
                return

            response = requests.get(image_attachments.url)
            message_attachment = Image.open(BytesIO(response.content))

            if message_attachment.width < message_attachment.height:
                message_attachment.thumbnail((200, message_attachment.height))
            else:
                message_attachment.thumbnail((message_attachment.width, 200))

            if style == "bubble":
                # The bubble is resized twice as for some reason .copy() doesn't work
                message_attachment.paste(
                    self._overlays["bubble"].resize(
                        (message_attachment.width, self._overlays["bubble"].height)
                    ),
                    (0, 0),
                    self._overlays["bubble"].resize(
                        (message_attachment.width, self._overlays["bubble"].height)
                    ),
                )
            elif style == "gang":
                message_attachment.paste(
                    self._overlays["gang"],
                    (
                        (
                            (message_attachment.width - self._overlays["gang"].width)
                            // 2
                        ),
                        (message_attachment.height - self._overlays["gang"].height),
                    ),
                    self._overlays["gang"],
                )

            message_file = BytesIO()
            message_attachment.save(message_file, format="PNG")
            message_file.seek(0)
            message_file = discord.File(message_file, filename="image.png")

            time_taken = datetime.datetime.now() - start_time
            embed = (
                discord.Embed(title=style.capitalize(), colour=discord.Colour.orange())
                .set_image(url="attachment://image.png")
                .set_footer(
                    text=f"{message_attachment.width}x{message_attachment.height}, "
                    f"{time_taken.microseconds / 1000}ms"
                )
            )

            await ctx.reply(embed=embed, file=message_file, mention_author=False)