import urllib.parse
import requests
import re
from io import BytesIO
from PIL import Image
import discord
from discord.ext import commands

cols = {
    "transparent": "",
    "white": "\\bg_white&space;",
    "black": "\\bg_black&space;",
    "red": "\\bg_red&space;",
    "green": "\\bg_green&space;",
    "blue": "\\bg_blue&space;",
}


class LatexCog(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @commands.command()
    async def latex(self, ctx, *, arguments):
        """Parses LaTeX and returns an image with the formatted result"""
        await send_latex(ctx, arguments)

    @commands.command()
    async def latexcol(self, ctx, *, arguments):
        """
        Parses LaTeX and returns an image with the formatted result
        using a specified color
        """
        res = re.search("\\s", arguments)
        if res is None:
            error = "Error: no color provided"
            error += "\nValid colors: %s" % (
                ", ".join("`%s`" % (x) for x in cols.keys())
            )
            await ctx.send(error)
            return
        split_ind = res.start()
        col = arguments[:split_ind].lower()
        inp = arguments[split_ind:].strip()
        await send_latex(ctx, inp, col)


async def send_latex(ctx, inp, col="black"):
    """Sends a response with the LaTeX output as an image.
    Args:
        ctx: The discord context
        inp: The LaTeX expression
        col (optional): The background color
    """
    latex = urllib.parse.quote(inp)
    try:
        colcode = cols[col]
    except KeyError:
        error = "Error: Invalid color: `%s`" % (col)
        error += "\nValid colors: %s" % (", ".join("`%s`" % (x) for x in cols.keys()))
        await ctx.send(error)
        return

    req_url = (
        "https://latex.codecogs.com/png.latex?\\dpi{300}&space;%s\\huge&space;%s"
        % (colcode, latex)
    )

    res = requests.get(req_url)
    img = Image.open(BytesIO(res.content))
    if col != "transparent":
        img = add_margin(img, 10, 10, 10, 10)

    buf = BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    disc_imgfile = discord.File(buf, filename="latex.png")

    embed = discord.Embed(title="LaTeX output", color=0x0539F5)
    embed.add_field(name="Input:", value="`%s`" % (inp), inline=False)
    embed.set_image(url="attachment://latex.png")
    embed.set_footer(text="Generated by codecogs.com")
    await ctx.send(embed=embed, file=disc_imgfile)


def add_margin(pil_img, top, right, bottom, left):
    """Adds margins to a PIL image

    This function relies on 2 preconditions:
        1. The image provided is using the palette ("P") color mode
        2. The first color in the palette is the background color
    Both of these preconditions should always hold for images from the codecogs API,
    but may not hold for all images, so this should not be used as a general purpose
    function for adding margins to an image.

    Args:
        pil_img: The PIL image to add margins to
        top, right, bottom, left: The width of the margin on each side

    Returns:
        A new PIL image with the margins added
    """
    width, height = pil_img.size
    new_width = width + right + left
    new_height = height + top + bottom
    result = Image.new("P", (new_width, new_height), 0)
    result.putpalette(pil_img.getpalette())
    result.paste(pil_img, (left, top))
    return result


def setup(bot):
    bot.add_cog(LatexCog(bot))
