path: root/bot/exts/avatar_modification/avatar_modify.py

import asyncio
import json
import logging
import math
import string
import unicodedata
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Callable, Optional, TypeVar, Union

import discord
from discord.ext import commands

from bot.bot import Bot
from bot.constants import Colours, Emojis
from bot.exts.avatar_modification._effects import PfpEffects
from bot.utils.extensions import invoke_help_command
from bot.utils.halloween import spookifications

log = logging.getLogger(__name__)

_EXECUTOR = ThreadPoolExecutor(10)

FILENAME_STRING = "{effect}_{author}.png"

MAX_SQUARES = 10_000

T = TypeVar("T")

GENDER_OPTIONS = json.loads(Path("bot/resources/holidays/pride/gender_options.json").read_text("utf8"))


async def in_executor(func: Callable[..., T], *args) -> T:
    """
    Runs the given synchronous function `func` in an executor.

    This is useful for running slow, blocking code within async
    functions, so that they don't block the bot.
    """
    log.trace(f"Running {func.__name__} in an executor.")
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(_EXECUTOR, func, *args)


def file_safe_name(effect: str, display_name: str) -> str:
    """Returns a file safe filename based on the given effect and display name."""
    valid_filename_chars = f"-_. {string.ascii_letters}{string.digits}"

    file_name = FILENAME_STRING.format(effect=effect, author=display_name)

    # Replace spaces
    file_name = file_name.replace(" ", "_")

    # Normalize unicode characters
    cleaned_filename = unicodedata.normalize("NFKD", file_name).encode("ASCII", "ignore").decode()

    # Remove invalid filename characters
    cleaned_filename = "".join(c for c in cleaned_filename if c in valid_filename_chars)
    return cleaned_filename


class AvatarModify(commands.Cog):
    """Various commands for users to apply affects to their own avatars."""

    def __init__(self, bot: Bot):
        self.bot = bot

    async def _fetch_user(self, user_id: int) -> Optional[discord.User]:
        """
        Fetches a user and handles errors.

        This helper function is required as the member cache doesn't always have the most up to date
        profile picture. This can lead to errors if the image is deleted from the Discord CDN.
        fetch_member can't be used due to the avatar url being part of the user object, and
        some weird caching that D.py does
        """
        try:
            user = await self.bot.fetch_user(user_id)
        except discord.errors.NotFound:
            log.debug(f"User {user_id} could not be found.")
            return None
        except discord.HTTPException:
            log.exception(f"Exception while trying to retrieve user {user_id} from Discord.")
            return None

        return user

    @commands.group(aliases=("avatar_mod", "pfp_mod", "avatarmod", "pfpmod"))
    async def avatar_modify(self, ctx: commands.Context) -> None:
        """Groups all of the pfp modifying commands to allow a single concurrency limit."""
        if not ctx.invoked_subcommand:
            await invoke_help_command(ctx)

    @avatar_modify.command(name="8bitify", root_aliases=("8bitify",))
    async def eightbit_command(self, ctx: commands.Context) -> None:
        """Pixelates your avatar and changes the palette to an 8bit one."""
        async with ctx.typing():
            user = await self._fetch_user(ctx.author.id)
            if not user:
                await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
                return

            image_bytes = await user.display_avatar.replace(size=1024).read()
            file_name = file_safe_name("eightbit_avatar", ctx.author.display_name)

            file = await in_executor(
                PfpEffects.apply_effect,
                image_bytes,
                PfpEffects.eight_bitify_effect,
                file_name
            )

            embed = discord.Embed(
                title="Your 8-bit avatar",
                description="Here is your avatar. I think it looks all cool and 'retro'."
            )

            embed.set_image(url=f"attachment://{file_name}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url)

        await ctx.send(embed=embed, file=file)

    @avatar_modify.command(name="reverse", root_aliases=("reverse",))
    async def reverse(self, ctx: commands.Context, *, text: Optional[str]) -> None:
        """
        Reverses the sent text.

        If no text is provided, the user's profile picture will be reversed.
        """
        if text:
            await ctx.send(f"> {text[::-1]}", allowed_mentions=discord.AllowedMentions.none())
            return

        async with ctx.typing():
            user = await self._fetch_user(ctx.author.id)
            if not user:
                await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
                return

            image_bytes = await user.display_avatar.replace(size=1024).read()
            filename = file_safe_name("reverse_avatar", ctx.author.display_name)

            file = await in_executor(
                PfpEffects.apply_effect,
                image_bytes,
                PfpEffects.flip_effect,
                filename
            )

            embed = discord.Embed(
                title="Your reversed avatar.",
                description="Here is your reversed avatar. I think it is a spitting image of you."
            )

            embed.set_image(url=f"attachment://{filename}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url)

            await ctx.send(embed=embed, file=file)

    @avatar_modify.command(aliases=("easterify",), root_aliases=("easterify", "avatareasterify"))
    async def avatareasterify(self, ctx: commands.Context, *colours: Union[discord.Colour, str]) -> None:
        """
        This "Easterifies" the user's avatar.

        Given colours will produce a personalised egg in the corner, similar to the egg_decorate command.
        If colours are not given, a nice little chocolate bunny will sit in the corner.
        Colours are split by spaces, unless you wrap the colour name in double quotes.
        Discord colour names, HTML colour names, XKCD colour names and hex values are accepted.
        """
        async def send(*args, **kwargs) -> str:
            """
            This replaces the original ctx.send.

            When invoking the egg decorating command, the egg itself doesn't print to to the channel.
            Returns the message content so that if any errors occur, the error message can be output.
            """
            if args:
                return args[0]

        async with ctx.typing():
            user = await self._fetch_user(ctx.author.id)
            if not user:
                await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
                return

            egg = None
            if colours:
                send_message = ctx.send
                ctx.send = send  # Assigns ctx.send to a fake send
                egg = await ctx.invoke(self.bot.get_command("eggdecorate"), *colours)
                if isinstance(egg, str):  # When an error message occurs in eggdecorate.
                    await send_message(egg)
                    return
                ctx.send = send_message  # Reassigns ctx.send

            image_bytes = await user.display_avatar.replace(size=256).read()
            file_name = file_safe_name("easterified_avatar", ctx.author.display_name)

            file = await in_executor(
                PfpEffects.apply_effect,
                image_bytes,
                PfpEffects.easterify_effect,
                file_name,
                egg
            )

            embed = discord.Embed(
                title="Your Lovely Easterified Avatar!",
                description="Here is your lovely avatar, all bright and colourful\nwith Easter pastel colours. Enjoy :D"
            )
            embed.set_image(url=f"attachment://{file_name}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=user.display_avatar.url)

        await ctx.send(file=file, embed=embed)

    @staticmethod
    async def send_pride_image(
        ctx: commands.Context,
        image_bytes: bytes,
        pixels: int,
        flag: str,
        option: str
    ) -> None:
        """Gets and sends the image in an embed. Used by the pride commands."""
        async with ctx.typing():
            file_name = file_safe_name("pride_avatar", ctx.author.display_name)

            file = await in_executor(
                PfpEffects.apply_effect,
                image_bytes,
                PfpEffects.pridify_effect,
                file_name,
                pixels,
                flag
            )

            embed = discord.Embed(
                title="Your Lovely Pride Avatar!",
                description=f"Here is your lovely avatar, surrounded by\n a beautiful {option} flag. Enjoy :D"
            )
            embed.set_image(url=f"attachment://{file_name}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=ctx.author.avatar.url)
            await ctx.send(file=file, embed=embed)

    @avatar_modify.group(
        aliases=("avatarpride", "pridepfp", "prideprofile"),
        root_aliases=("prideavatar", "avatarpride", "pridepfp", "prideprofile"),
        invoke_without_command=True
    )
    async def prideavatar(self, ctx: commands.Context, option: str = "lgbt", pixels: int = 64) -> None:
        """
        This surrounds an avatar with a border of a specified LGBT flag.

        This defaults to the LGBT rainbow flag if none is given.
        The amount of pixels can be given which determines the thickness of the flag border.
        This has a maximum of 512px and defaults to a 64px border.
        The full image is 1024x1024.
        """
        option = option.lower()
        pixels = max(0, min(512, pixels))
        flag = GENDER_OPTIONS.get(option)
        if flag is None:
            await ctx.send("I don't have that flag!")
            return

        async with ctx.typing():
            user = await self._fetch_user(ctx.author.id)
            if not user:
                await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
                return
            image_bytes = await user.display_avatar.replace(size=1024).read()
            await self.send_pride_image(ctx, image_bytes, pixels, flag, option)

    @prideavatar.command()
    async def flags(self, ctx: commands.Context) -> None:
        """This lists the flags that can be used with the prideavatar command."""
        choices = sorted(set(GENDER_OPTIONS.values()))
        options = "• " + "\n• ".join(choices)
        embed = discord.Embed(
            title="I have the following flags:",
            description=options,
            colour=Colours.soft_red
        )
        await ctx.send(embed=embed)

    @avatar_modify.command(
        aliases=("savatar", "spookify"),
        root_aliases=("spookyavatar", "spookify", "savatar"),
        brief="Spookify an user's avatar."
    )
    async def spookyavatar(self, ctx: commands.Context) -> None:
        """This "spookifies" the user's avatar, with a random *spooky* effect."""
        user = await self._fetch_user(ctx.author.id)
        if not user:
            await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
            return

        async with ctx.typing():
            image_bytes = await user.display_avatar.replace(size=1024).read()

            file_name = file_safe_name("spooky_avatar", ctx.author.display_name)

            file = await in_executor(
                PfpEffects.apply_effect,
                image_bytes,
                spookifications.get_random_effect,
                file_name
            )

            embed = discord.Embed(
                title="Is this you or am I just really paranoid?",
                colour=Colours.soft_red
            )
            embed.set_image(url=f"attachment://{file_name}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}.", icon_url=ctx.author.display_avatar.url)

            await ctx.send(file=file, embed=embed)

    @avatar_modify.command(name="mosaic", root_aliases=("mosaic",))
    async def mosaic_command(self, ctx: commands.Context, squares: int = 16) -> None:
        """Splits your avatar into x squares, randomizes them and stitches them back into a new image!"""
        async with ctx.typing():
            user = await self._fetch_user(ctx.author.id)
            if not user:
                await ctx.send(f"{Emojis.cross_mark} Could not get user info.")
                return

            if not 1 <= squares <= MAX_SQUARES:
                raise commands.BadArgument(f"Squares must be a positive number less than or equal to {MAX_SQUARES:,}.")

            sqrt = math.sqrt(squares)

            if not sqrt.is_integer():
                squares = math.ceil(sqrt) ** 2  # Get the next perfect square

            file_name = file_safe_name("mosaic_avatar", ctx.author.display_name)

            img_bytes = await user.display_avatar.replace(size=1024).read()

            file = await in_executor(
                PfpEffects.apply_effect,
                img_bytes,
                PfpEffects.mosaic_effect,
                file_name,
                squares,
            )

            if squares == 1:
                title = "Hooh... that was a lot of work"
                description = "I present to you... Yourself!"
            elif squares == MAX_SQUARES:
                title = "Testing the limits I see..."
                description = "What a masterpiece. :star:"
            else:
                title = "Your mosaic avatar"
                description = f"Here is your avatar. I think it looks a bit *puzzling*\nMade with {squares} squares."

            embed = discord.Embed(
                title=title,
                description=description,
                colour=Colours.blue
            )

            embed.set_image(url=f"attachment://{file_name}")
            embed.set_footer(text=f"Made by {ctx.author.display_name}", icon_url=user.display_avatar.url)

            await ctx.send(file=file, embed=embed)


def setup(bot: Bot) -> None:
    """Load the AvatarModify cog."""
    bot.add_cog(AvatarModify(bot))