path: root/bot/exts/utilities/wolfram.py

from collections.abc import Callable
from io import BytesIO
from urllib.parse import urlencode

import arrow
import discord
from discord import Embed
from discord.ext import commands
from discord.ext.commands import BucketType, Cog, Context, check, group
from pydis_core.utils.logging import get_logger

from bot.bot import Bot
from bot.constants import Colours, STAFF_ROLES, Wolfram as WolframConfig
from bot.utils.pagination import ImagePaginator

log = get_logger(__name__)

APPID = WolframConfig.key.get_secret_value()
DEFAULT_OUTPUT_FORMAT = "JSON"
QUERY = "http://api.wolframalpha.com/v2/{request}"
WOLF_IMAGE = "https://www.symbols.com/gi.php?type=1&id=2886&i=1"

MAX_PODS = 20

# Allows for 10 wolfram calls pr user pr day
usercd = commands.CooldownMapping.from_cooldown(WolframConfig.user_limit_day, 60 * 60 * 24, BucketType.user)

# Allows for max api requests / days in month per day for the entire guild (Temporary)
guildcd = commands.CooldownMapping.from_cooldown(WolframConfig.guild_limit_day, 60 * 60 * 24, BucketType.guild)


async def send_embed(
        ctx: Context,
        message_txt: str,
        colour: int = Colours.soft_red,
        footer: str | None = None,
        img_url: str | None = None,
        f: discord.File = None
) -> None:
    """Generate & send a response embed with Wolfram as the author."""
    embed = Embed(colour=colour)
    embed.description = message_txt
    embed.set_author(
        name="Wolfram Alpha",
        icon_url=WOLF_IMAGE,
        url="https://www.wolframalpha.com/"
    )
    if footer:
        embed.set_footer(text=footer)

    if img_url:
        embed.set_image(url=img_url)

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


def custom_cooldown(*ignore: int) -> Callable:
    """
    Implement per-user and per-guild cooldowns for requests to the Wolfram API.

    A list of roles may be provided to ignore the per-user cooldown.
    """
    async def predicate(ctx: Context) -> bool:
        if ctx.invoked_with == "help":
            # if the invoked command is help we don't want to increase the ratelimits since it's not actually
            # invoking the command/making a request, so instead just check if the user/guild are on cooldown.
            guild_cooldown = guildcd.get_bucket(ctx.message).get_tokens() != 0  # if guild is on cooldown
            # check the message is in a guild, and check user bucket if user is not ignored
            if ctx.guild and not any(r.id in ignore for r in ctx.author.roles):
                return guild_cooldown and usercd.get_bucket(ctx.message).get_tokens() != 0
            return guild_cooldown

        user_bucket = usercd.get_bucket(ctx.message)

        if all(role.id not in ignore for role in ctx.author.roles):
            user_rate = user_bucket.update_rate_limit()

            if user_rate:
                # Can't use api; cause: member limit
                cooldown = arrow.utcnow().shift(seconds=int(user_rate)).humanize(only_distance=True)
                message = (
                    "You've used up your limit for Wolfram|Alpha requests.\n"
                    f"Cooldown: {cooldown}"
                )
                await send_embed(ctx, message)
                return False

        guild_bucket = guildcd.get_bucket(ctx.message)
        guild_rate = guild_bucket.update_rate_limit()

        # Repr has a token attribute to read requests left
        log.debug(guild_bucket)

        if guild_rate:
            # Can't use api; cause: guild limit
            message = (
                "The max limit of requests for the server has been reached for today.\n"
                f"Cooldown: {int(guild_rate)}"
            )
            await send_embed(ctx, message)
            return False

        return True

    return check(predicate)


async def get_pod_pages(ctx: Context, bot: Bot, query: str) -> list[tuple[str, str]] | None:
    """Get the Wolfram API pod pages for the provided query."""
    async with ctx.typing():
        params = {
            "input": query,
            "appid": APPID,
            "output": DEFAULT_OUTPUT_FORMAT,
            "format": "image,plaintext",
            "location": "the moon",
            "latlong": "0.0,0.0",
            "ip": "1.1.1.1"
        }
        request_url = QUERY.format(request="query")

        async with bot.http_session.get(url=request_url, params=params) as response:
            json = await response.json(content_type="text/plain")

        result = json["queryresult"]
        log_full_url = f"{request_url}?{urlencode(params)}"
        if result["error"]:
            # API key not set up correctly
            if result["error"]["msg"] == "Invalid appid":
                message = "Wolfram API key is invalid or missing."
                log.warning(
                    "API key seems to be missing, or invalid when "
                    f"processing a wolfram request: {log_full_url}, Response: {json}"
                )
                await send_embed(ctx, message)
                return None

            message = "Something went wrong internally with your request, please notify staff!"
            log.warning(f"Something went wrong getting a response from wolfram: {log_full_url}, Response: {json}")
            await send_embed(ctx, message)
            return None

        if not result["success"]:
            message = f"I couldn't find anything for {query}."
            await send_embed(ctx, message)
            return None

        if not result["numpods"]:
            message = "Could not find any results."
            await send_embed(ctx, message)
            return None

        pods = result["pods"]
        pages = []
        for pod in pods[:MAX_PODS]:
            subs = pod.get("subpods")

            for sub in subs:
                title = sub.get("title") or sub.get("plaintext") or sub.get("id", "")
                img = sub["img"]["src"]
                pages.append((title, img))
        return pages


class Wolfram(Cog):
    """Commands for interacting with the Wolfram|Alpha API."""

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

    @group(name="wolfram", aliases=("wolf", "wa"), invoke_without_command=True)
    @custom_cooldown(*STAFF_ROLES)
    async def wolfram_command(self, ctx: Context, *, query: str) -> None:
        """Requests all answers on a single image, sends an image of all related pods."""
        params = {
            "i": query,
            "appid": APPID,
            "location": "the moon",
            "latlong": "0.0,0.0",
            "ip": "1.1.1.1"
        }
        request_url = QUERY.format(request="simple")

        # Give feedback that the bot is working.
        async with ctx.typing():
            async with self.bot.http_session.get(url=request_url, params=params) as response:
                status = response.status
                image_bytes = await response.read()

            f = discord.File(BytesIO(image_bytes), filename="image.png")
            image_url = "attachment://image.png"

            if status == 501:
                message = "Failed to get response."
                footer = ""
                color = Colours.soft_red
            elif status == 400:
                message = "No input found."
                footer = ""
                color = Colours.soft_red
            elif status == 403:
                message = "Wolfram API key is invalid or missing."
                footer = ""
                color = Colours.soft_red
            elif status != 200:
                # Handle all other possible status codes here
                message = f"Unexpected status code from Wolfram API: {status}"
                footer = ""
                color = Colours.soft_red

                log.warning(f"Unexpected status code from Wolfram API: {status}\nInput: {query}")
            else:
                message = ""
                footer = "View original for a bigger picture."
                color = Colours.soft_orange

            # Sends a "blank" embed if no request is received, unsure how to fix
            await send_embed(ctx, message, color, footer=footer, img_url=image_url, f=f)

    @wolfram_command.command(name="page", aliases=("pa", "p"))
    @custom_cooldown(*STAFF_ROLES)
    async def wolfram_page_command(self, ctx: Context, *, query: str) -> None:
        """
        Requests a drawn image of given query.

        Keywords worth noting are, "like curve", "curve", "graph", "pokemon", etc.
        """
        pages = await get_pod_pages(ctx, self.bot, query)

        if not pages:
            return

        embed = Embed()
        embed.set_author(
            name="Wolfram Alpha",
            icon_url=WOLF_IMAGE,
            url="https://www.wolframalpha.com/"
        )
        embed.colour = Colours.soft_orange

        await ImagePaginator.paginate(pages, ctx, embed)

    @wolfram_command.command(name="cut", aliases=("c",))
    @custom_cooldown(*STAFF_ROLES)
    async def wolfram_cut_command(self, ctx: Context, *, query: str) -> None:
        """
        Requests a drawn image of given query.

        Keywords worth noting are, "like curve", "curve", "graph", "pokemon", etc.
        """
        pages = await get_pod_pages(ctx, self.bot, query)

        if not pages:
            return

        page = pages[1] if len(pages) >= 2 else pages[0]

        await send_embed(ctx, page[0], colour=Colours.soft_orange, img_url=page[1])

    @wolfram_command.command(name="short", aliases=("sh", "s"))
    @custom_cooldown(*STAFF_ROLES)
    async def wolfram_short_command(self, ctx: Context, *, query: str) -> None:
        """Requests an answer to a simple question."""
        params = {
            "i": query,
            "appid": APPID,
            "location": "the moon",
            "latlong": "0.0,0.0",
            "ip": "1.1.1.1"
        }
        request_url = QUERY.format(request="result")

        # Give feedback that the bot is working.
        async with ctx.typing():
            async with self.bot.http_session.get(url=request_url, params=params) as response:
                status = response.status
                response_text = await response.text()

            if status == 501:
                message = "Failed to get response."
                color = Colours.soft_red
            elif status == 400:
                message = "No input found."
                color = Colours.soft_red
            elif response_text == "Error 1: Invalid appid.":
                message = "Wolfram API key is invalid or missing."
                color = Colours.soft_red
            elif status != 200:
                # Handle all other possible status codes here
                message = f"Unexpected status code from Wolfram API: {status}"
                color = Colours.soft_red

                log.warning(f"Unexpected status code from Wolfram API: {status}\nInput: {query}")
            else:
                message = response_text
                color = Colours.soft_orange

            await send_embed(ctx, message, color)


async def setup(bot: Bot) -> None:
    """Load the Wolfram cog."""
    if not WolframConfig.key:
        log.warning("No Wolfram API Key was provided. Not loading Wolfram Cog.")
        return
    await bot.add_cog(Wolfram(bot))