modufur/commands/booru.py

import urlextract
import hikari
import lightbulb
import pysaucenao

from tools import components, scraper


plugin = lightbulb.Plugin("booru")
extractor = urlextract.URLExtract()


@plugin.command
# @lightbulb.option('attachment', 'Attachment(s) to reverse')
@lightbulb.option("url", "URL(s) to reverse, separated by space")
@lightbulb.command("reverse", "Reverse image search using SauceNAO & Kheina", ephemeral=True)
@lightbulb.implements(lightbulb.SlashCommand, lightbulb.MessageCommand)
async def reverse(context):
    match context:
        case lightbulb.SlashContext():
            urls = extractor.find_urls(context.options.url or "", only_unique=True, with_schema_only=True)

            if not urls:
                await context.respond("**Invalid URL(s).**")
                return

            await _reverse(context, urls)
        case lightbulb.MessageContext():
            urls = extractor.find_urls(context.options.target.content or "", only_unique=True, with_schema_only=True)
            urls += [attachment.url for attachment in context.options.target.attachments if attachment.url not in urls]

            if not urls:
                await context.respond("**No images found.**")
                return

            selector = None

            if len(urls) > 1:
                selector = components.Selector(
                    pages=[
                        f"**Select potential images to search: `{urls.index(url) + 1}/{len(urls)}`**\n{url}"
                        for url in urls
                    ],
                    buttons=[components.Back(), components.Forward(), components.Select(), components.Confirm()],
                    urls=urls,
                )

                await selector.send(context.interaction, ephemeral=True)
                await selector.wait()

                if selector.timed_out:
                    await context.interaction.edit_initial_response("**Timed out.**", components=None)
                    return

                urls = selector.selected

            await _reverse(context, urls, selector=selector)


@reverse.set_error_handler()
async def on_reverse_error(event):
    error = None

    match event.exception.__cause__:
        case pysaucenao.ShortLimitReachedException():
            error = "**API limit reached. Please try again in a minute.**"
        case pysaucenao.DailyLimitReachedException():
            error = "**Daily API limit reached. Please try again tomorrow.**"
        case pysaucenao.FileSizeLimitException() as url:
            error = f"**Image file size too large.**\n{url}"
        case pysaucenao.ImageSizeException() as url:
            error = f"**Image resolution too small.**\n{url}"
        case pysaucenao.InvalidImageException() as url:
            error = f"**Invalid image.**\n{url}"
        case pysaucenao.UnknownStatusCodeException():
            error = "**An unknown SauceNAO error has occurred. The service may be down.**"

    if error:
        try:
            await event.context.respond(error)
        except:
            await event.context.interaction.edit_initial_response(error, components=None)

        return True


async def _reverse(context, urls, *, selector=None):
    if not selector:
        await context.respond(hikari.ResponseType.DEFERRED_MESSAGE_CREATE)

    matches = await scraper.reverse(urls)

    if not matches:
        if selector:
            await context.interaction.edit_initial_response("**No matches found.**", components=None)
        else:
            await context.respond("**No matches found.**")
        return

    pages = [
        (
            hikari.Embed(
                title=match["artist"], url=match["url"], color=context.get_guild().get_my_member().get_top_role().color
            )
            .set_author(name=f'{match["similarity"]}% Match')
            .set_image(match["thumbnail"])
            .set_footer(match["source"])
        )
        if match
        else f"**No match found.**\n{urls[index]}"
        for index, match in enumerate(matches)
    ]

    if len(pages) > 1:
        selector = components.Selector(pages=pages, buttons=[components.Back(), components.Forward()], timeout=900)

        await selector.send_edit(context.interaction)
    else:
        if selector:
            await context.interaction.edit_initial_response(content=None, embed=pages[0], components=None)
        else:
            await context.respond(pages[0])


def load(bot):
    bot.add_plugin(plugin)


def unload(bot):
    bot.remove_plugin(plugin)