#!/usr/bin/env python3
import json
import zlib
import sys
import pathlib
import time
import urllib.request
import shutil

# PNG, APNG, Lottie, GIF
STICKER_EXTENSIONS = ["", "png", "png", "json", "gif"]


def download_file(url, path):
    if path.exists():
        print("Already downloaded, skipping...")
        return
    out_file = open(path, "wb")
    try:
        req = urllib.request.urlopen(
            urllib.request.Request(url, headers={"User-Agent": "curl/8.6.0"})
        )
        shutil.copyfileobj(req, out_file)
    except urllib.error.HTTPError:
        print(f"Failed to download {url}")
    time.sleep(0.1)


def get_guilds(filename):
    if filename.endswith(".json"):
        # TODO: specify if it's compressed or not better
        jsondata = open(filename, "r").read()
    else:
        bindata = open(filename, "rb").read()
        decompress = zlib.decompressobj()
        jsondata = decompress.decompress(bindata).decode("utf-8")

    decoder = json.JSONDecoder()
    guilds = []
    while (offset := jsondata.find("{")) != -1:
        obj, end = decoder.raw_decode(jsondata[offset:])
        jsondata = jsondata[end + offset :]

        typ = obj.get("t", "")
        if typ == "READY":
            guilds.extend(obj["d"]["guilds"])
        elif typ == "GUILD_CREATE":
            guilds.append(obj["d"])

    return guilds


def safe_name(s):
    # if you're on windows do something else here, i don't care
    return s.replace("/", " ")


def main(*args):
    guilds = get_guilds(args[1])

    base_dir = pathlib.Path("out") / "discord"

    for guild in guilds:
        name = guild["properties"]["name"]
        print(f"Processing guild '{name}'")

        sticker_objs = guild["stickers"]
        emoji_objs = guild["emojis"]

        dir_name = safe_name(name)
        guild_dir = base_dir / dir_name
        sticker_dir = guild_dir / "stickers"
        emoji_dir = guild_dir / "emoji"

        guild_dir.mkdir(exist_ok=True, parents=True)

        properties = guild["properties"]
        guild_id = properties["id"]
        subset_properties = {
            "id": properties["id"],
            "name": properties["name"],
            "description": properties["description"],
        }
        with open(guild_dir / "properties.json", "w") as properties_file:
            json.dump(subset_properties, properties_file)

        if properties["icon"] is not None:
            download_file(
                f"https://cdn.discordapp.com/icons/{guild_id}/{properties['icon']}.png?size=4096",
                guild_dir / "icon.png",
            )
        if properties["banner"] is not None:
            download_file(
                f"https://cdn.discordapp.com/banners/{guild_id}/{properties['banner']}.png?size=4096",
                guild_dir / "banner.png",
            )

        if len(sticker_objs) > 0:
            sticker_dir.mkdir(exist_ok=True, parents=True)
        if len(emoji_objs) > 0:
            emoji_dir.mkdir(exist_ok=True, parents=True)

        for sticker in sticker_objs:
            print(f"Found sticker '{sticker['name']}'")
            extension = STICKER_EXTENSIONS[sticker["format_type"]]
            if extension == "json":
                # Lottie is weird for some reason
                url = f"https://discord.com/stickers/{sticker['id']}.json"
            else:
                url = f"https://media.discordapp.net/stickers/{sticker['id']}.{extension}?size=4096"
            # Set passthrough=false for APNGs if you don't want animated

            filename = sticker_dir / f"{safe_name(sticker['name'])}.{extension}"
            print(f"Downloading {url} to {filename}")
            download_file(url, filename)

        for emoji in emoji_objs:
            print(f"Found emoji '{emoji['name']}'")
            extension = "gif" if emoji["animated"] else "png"
            url = f"https://cdn.discordapp.com/emojis/{emoji['id']}.{extension}"
            filename = emoji_dir / f"{safe_name(emoji['name'])}.{extension}"
            print(f"Downloading {url} to {filename}")
            download_file(url, filename)


if __name__ == "__main__":
    main(*sys.argv)