<filename>main.py<gh_stars>0
import shutil
import asyncio
import requests
import os
import discord
from dotenv import load_dotenv

load_dotenv()
TOKEN = os.getenv('DISCORD_TOKEN')
CMD_PREFIX = '!'

wiimmInProgress = False

checkpoint_maps = {
    'lc': 'beginner_course.kmp.png',
    'mmm': 'farm_course.kmp.png',
    'mg': 'kinoko_course.kmp.png',
    'tf': 'factory_course.kmp.png',
    'mc': 'castle_course.kmp.png',
    'cm': 'shopping_course.kmp.png',
    'dks': 'boardcross_course.kmp.png',
    'dksc': 'boardcross_course.kmp.png',
    'wgm': 'truck_course.kmp.png',
    'dc': 'senior_course.kmp.png',
    'kc': 'water_course.kmp.png',
    'mt': 'treehouse_course.kmp.png',
    'gv': 'volcano_course.kmp.png',
    'ddr': 'desert_course.kmp.png',
    'mh': 'ridgehighway_course.kmp.png',
    'bc': 'koopa_course.kmp.png',
    'rr': 'rainbow_course.kmp.png',
    'rpb': 'old_peach_gc.kmp.png',
    'ryf': 'old_falls_ds.kmp.png',
    'rgv2': 'old_obake_sfc.kmp.png',
    'rmr': 'old_mario_64.kmp.png',
    'rsl': 'old_sherbet_64.kmp.png',
    'rsgb': 'old_heyho_gba.kmp.png',
    'rds': 'old_town_ds.kmp.png',
    'rws': 'old_waluigi_gc.kmp.png',
    'rdh': 'old_desert_ds.kmp.png',
    'rbc3': 'old_koopa_gba.kmp.png',
    'rdkjp': 'old_donkey_64.kmp.png',
    'rmc': 'old_mario_gc.kmp.png',
    'rmc3': 'old_mario_sfc.kmp.png',
    'rpg': 'old_garden_ds.kmp.png',
    'rdkm': 'old_donkey_gc.kmp.png',
    'rbc': 'old_koopa_64.kmp.png'
}


class Bot(discord.Client):
    async def on_ready(self):
        print(f"Bot connected as {self.user}.")

    async def cmdHelp(self, cmd):
        response = ''
        if cmd in 'cp':
            response += f"```css\n!cp <track>\n```"
        if cmd in 'id':
            response += f"```css\n!id <ct.wiimm.de track ID>\n```"
        return response

    async def embedError(self, msg, error, cmd, details):
        cmdHelp = await self.cmdHelp(cmd)
        if 'track' in details:
            cmdHelp += "\nUse one of the following:\n" + await self.printTracks()
        if 'notint' in details:
            cmdHelp += "\nMake sure you aren't including any letters in your search."
        if 'outofrange' in details:
            cmdHelp += "\nThe ID provided is not associated with a custom track or arena."
        if 'convert' in details:
            cmdHelp += "\nThis track could not be converted to a proper format. This is likely due to the track having an \'invalid\' class."
        if 'generate' in details:
            cmdHelp += "\nThis checkpoint map could not be generated. This is likely due to the track having an \'invalid\' class."
        embed = discord.Embed(title=error, description=cmdHelp, color=0xC13353)
        return await msg.channel.send(embed=embed)

    async def on_message(self, msg):
        # Prevent recursion
        if msg.author == self.user:
            return

        # Check for bot command
        if len(msg.content) > 0 and msg.content[0] == CMD_PREFIX:
            await self.parseBotCmd(msg)

    async def parseBotCmd(self, msg):
        msgContent = msg.content[1:].split()
        cmd = msgContent[0].lower()

        if cmd == 'help':
            return await self.listCommands(msg)
        if cmd == 'cp':
            return await self.cp(msg, msgContent[1:])
        if cmd == 'id':
            if not wiimmInProgress:
                return await self.customCheckpointMaps(msg, msgContent[1:])
            else:
                embed = discord.Embed(title="Command already in use", color=0xC13353)
                return await msg.channel.send(embed=embed)

    async def listCommands(self, msg):
        description = "**\*** !cp <track> - Sends the checkpoint map for selected Nintendo track\n" \
                      "**\*** !id <ct.wiimm.de id> - Generates/sends the checkpoint map for the selected custom track"
        embed = discord.Embed(title="List of available commands", description=description, color=0x89DA72)
        return await msg.channel.send(embed=embed)

    async def cp(self, msg, msgContent):
        try:
            track = msgContent[0]
        except IndexError:
            error = "No track provided"
            details = 'track'
            cmd = 'cp'
            await self.embedError(msg, error, cmd, details)
            return

        if len(track) > 0:
            try:
                temp = checkpoint_maps[track.lower()]
                return await msg.channel.send(file=discord.File("./resources/maps/" + temp))
            except:
                error = "Invalid track name"
                details = 'track'
                cmd = 'cp'
                await self.embedError(msg, error, cmd, details)
                return

    async def customCheckpointMaps(self, msg, msgContent):
        global wiimmInProgress

        try:
            track_id = int(msgContent[0])
        except:
            error = "ID cannot exist"
            details = 'notint'
            cmd = 'id'
            await self.embedError(msg, error, cmd, details)
            return

        track_id = str(track_id)
        if os.path.isfile(f"./resources/custom_maps/{track_id}.png"):
            return await msg.channel.send(file=discord.File("./resources/custom_maps/"+track_id+".png"))
        else:
            wiimmInProgress = True

            if os.path.isfile(f"./resources/custom_tracks/{track_id}.szs"):
                return await self.generateMap(msg, track_id)
            elif os.path.isfile(f"./resources/downloads/{track_id}.wbz"):
                return await self.convertWbz(msg, track_id)
            else:
                return await self.downloadWbz(msg, track_id)

    async def downloadWbz(self, msg, track_id):
        global wiimmInProgress

        try:
            embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                                  description=f"Downloading new ID ({track_id})...",
                                  color=0x7289DA)
            progress_msg = await msg.channel.send(embed=embed)
        except discord.HTTPException as e:
            if e.code == 50035:
                wiimmInProgress = False
                error = "Track does not exist"
                details = 'outofrange'
                cmd = 'id'
                await self.embedError(msg, error, cmd, details)
                return

        url = f"https://ct.wiimm.de/d/{track_id}"
        r = requests.get(url, allow_redirects=True)

        with open(f"./resources/downloads/{track_id}.wbz", 'wb') as f:
            f.write(r.content)

        if r.headers['content-type'] != 'application/octet-stream':
            os.remove(f"./resources/downloads/{track_id}.wbz")
            wiimmInProgress = False
            error = "Track does not exist"
            details = 'outofrange'
            cmd = 'id'
            await self.embedError(msg, error, cmd, details)
            return

        embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                              description=f"{track_id}.wbz downloaded successfully",
                              color=0x7289DA)
        await progress_msg.edit(embed=embed)
        return await self.convertWbz(msg, track_id, progress_msg)

    async def convertWbz(self, msg, track_id, progress_msg=None):
        global wiimmInProgress

        if progress_msg is None:
            embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                                  description="Converting to SZS...",
                                  color=0x7289DA)
            progress_msg = await msg.channel.send(embed=embed)
        else:
            embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                                  description="Converting to SZS...",
                                  color=0x7289DA)
            await progress_msg.edit(embed=embed)

        if not os.path.isdir("./wbz-to-szs"):
            os.mkdir("./wbz-to-szs")

        shutil.copy("./resources/downloads/" + track_id + ".wbz", "./wbz-to-szs/temp.wbz")
        convert = f"wszst normalize ./wbz-to-szs/temp.wbz --autoadd-path ./auto-add --DEST ./WBZ-to-SZS/temp.szs --szs --overwrite"

        c_proc = await asyncio.create_subprocess_shell(convert)
        await c_proc.communicate()

        try:
            shutil.copy("./wbz-to-szs/temp.szs", "./resources/custom_tracks/" + track_id + ".szs")
        except FileNotFoundError:
            wiimmInProgress = bool(0)
            error = "Track could not be converted"
            details = 'convert'
            cmd = 'id'
            await self.embedError(msg, error, cmd, details)
            return

        os.remove(f"./wbz-to-szs/temp.wbz")
        os.remove(f"./wbz-to-szs/temp.szs")

        embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                              description=f"{track_id}.szs converted successfully",
                              color=0x7289DA)
        await progress_msg.edit(embed=embed)
        return await self.generateMap(msg, track_id, progress_msg)

    async def generateMap(self, msg, track_id, progress_msg=None):
        global wiimmInProgress

        if progress_msg is None:
            embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                                  description="Generating checkpoint map...",
                                  color=0x7289DA)
            progress_msg = await msg.channel.send(embed=embed)
        else:
            embed = discord.Embed(title=f"Creating new checkpoint map ({track_id})...",
                                  description="Generating checkpoint map...",
                                  color=0x7289DA)
            await progress_msg.edit(embed=embed)

        draw = "wkmpt draw ./resources/custom_tracks/" + track_id + ".szs --kcl-script road.txt --draw=ckpt,kcl,roadobjects,black --png 100a0"

        g_proc = await asyncio.create_subprocess_shell(draw)
        await g_proc.communicate()

        try:
            shutil.copy("./resources/custom_tracks/" + track_id + ".kmp.png",
                        "./resources/custom_maps/" + track_id + ".png")
        except FileNotFoundError:
            wiimmInProgress = False
            error = "Checkpoint map could not be generated"
            details = 'generate'
            cmd = 'id'
            await self.embedError(msg, error, cmd, details)
            return

        os.remove(f"./resources/custom_tracks/{track_id}.kmp.png")

        author = str(msg.author.id)

        file = discord.File(f"./resources/custom_maps/{track_id}.png", filename=f"{track_id}.png")
        embed = discord.Embed(title=f"Successfully created checkpoint map ({track_id})",
                              description=f"New CT Archive checkpoint map created! <@{author}>",
                              color=0x89DA72)
        embed.set_image(url=f"attachment://{track_id}.png")
        await progress_msg.delete()
        await msg.channel.send(file=file, embed=embed)

        wiimmInProgress = False

        return

    async def printTracks(self):
        response = "```LC     MC     DC     DDR\n" \
                   "MMM    CM     KC     MH\n" \
                   "MG     DKS    MT     BC\n" \
                   "TF     WGM    GV     RR\n\n" \
                   "rPB    rSL    rDH    rMC3\n" \
                   "rYF    rSGB   rBC3   rPG\n" \
                   "rGV2   rDS    rDKJP  rDKM\n" \
                   "rMR    rWS    rMC    rBC```"
        return response


client = Bot()
client.run(TOKEN)
