homelab/modules/spin.py

import asyncio
import random
from io import BytesIO

import aiohttp

# noinspection PyUnresolvedReferences
from modules.squotes import render_message
from pyrogram import Client, enums, filters, types
from pyrogram.types import Message

# noinspection PyUnresolvedReferences
from utils.misc import modules_help, prefix
from utils.scripts import format_exc, import_library, resize_image

Image = import_library("PIL", "pillow").Image
np = import_library("numpy")
imageio = import_library("imageio")


def create_gif(filename: str, offset: int, fps: int = 2, typ: str = "spin"):
    img = Image.open(f"downloads/{filename}")
    if typ.lower() != "spin":
        img = img.resize(
            (random.randint(200, 1280), random.randint(200, 1280)), Image.ANTIALIAS
        )
    imageio.mimsave(
        "downloads/video.gif",
        [img.rotate(-(i % 360)) for i in range(1, 361, offset)],
        fps=fps,
    )


async def quote_cmd(client: Client, message: types.Message):
    count = 1

    is_png = False
    send_for_me = False
    no_reply = False

    messages = []

    async for msg in client.get_chat_history(
        message.chat.id, offset_id=message.reply_to_message.id, reverse=True
    ):
        if msg.empty:
            continue
        if msg.message_id >= message.id:
            break
        if no_reply:
            msg.reply_to_message = None

        messages.append(msg)

        if len(messages) >= count:
            break

    if send_for_me:
        await message.delete()
        message = await client.send_message(
            "me", "<b>Generating...</b>", parse_mode=enums.ParseMode.HTML
        )
    else:
        await message.edit("<b>Generating...</b>", parse_mode=enums.ParseMode.HTML)

    url = "https://quotes.fl1yd.su/generate"
    params = {
        "messages": [
            await render_message(client, msg) for msg in messages if not msg.empty
        ],
        "quote_color": "#162330",
        "text_color": "#fff",
    }

    response = await aiohttp.ClientSession().post(url, json=params)
    if response.status != 200:
        return await message.edit(
            f"<b>Quotes API error!</b>\n" f"<code>{response.text}</code>",
            parse_mode=enums.ParseMode.HTML,
        )

    resized = resize_image(
        BytesIO(await response.read()), img_type="PNG" if is_png else "WEBP"
    )
    return resized, is_png


@Client.on_message(filters.command(["spin", "dspin"], prefix) & filters.me)
async def spin_handler(client: Client, message: Message):
    if not message.reply_to_message:
        await message.edit(
            "<b>Reply to a <i>message</i> to spin it!</b>",
            parse_mode=enums.ParseMode.HTML,
        )
        await message.edit(
            "<b>Downloading sticker...</b>", parse_mode=enums.ParseMode.HTML
        )
        return await message.edit(
            "<b>Invalid file type!</b>", parse_mode=enums.ParseMode.HTML
        )
    try:
        coro = True
        if message.reply_to_message.document:
            filename = message.reply_to_message.document.file_name
            if (
                not filename.endswith(".webp")
                and not filename.endswith(".png")
                and not filename.endswith(".jpg")
                and not filename.endswith(".jpeg")
            ):
                return await message.edit(
                    "<b>Invalid file type!</b>", parse_mode=enums.ParseMode.HTML
                )
        elif message.reply_to_message.sticker:
            if message.reply_to_message.sticker.is_video:
                return await message.edit(
                    "<b>Video stickers not allowed</b>", parse_mode=enums.ParseMode.HTML
                )
            filename = "sticker.webp"
        elif message.reply_to_message.text:
            result = await quote_cmd(client, message)
            if result[1]:
                filename = "sticker.png"
            else:
                filename = "sticker.webp"
            open(f"downloads/" + filename, "wb").write(result[0].getbuffer())
            coro = False
        else:
            filename = "photo.jpg"
        if coro:
            await message.reply_to_message.download(f"downloads/{filename}")
    except Exception as ex:
        return await message.edit(
            f"<b>Message can not be loaded:</b>\n<code>{format_exc(ex)}</code>",
            parse_mode=enums.ParseMode.HTML,
        )
    await message.edit("<b>Spinning...</b>", parse_mode=enums.ParseMode.HTML)
    offset = int(message.command[1]) if len(message.command) > 1 else 10
    fps = int(message.command[2]) if len(message.command) > 2 else 30
    try:
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(
            None, lambda: create_gif(filename, offset, fps, message.command[0])
        )
        await message.delete()
        return await client.send_animation(
            chat_id=message.chat.id,
            animation="downloads/video.gif",
            reply_to_message_id=message.reply_to_message.id,
        )
    except Exception as e:
        await message.reply(format_exc(e), parse_mode=enums.ParseMode.HTML)


modules_help["spin"] = {
    "spin [offset] [fps]": "Spin message (Reply required)",
    "dspin [offset] [fps]": "SHAKAL spin message (Reply required)",
}