from asyncio import run, sleep
import json
from aiohttp import ClientSession
import aiofiles
from pathlib import Path
from yarl import URL
from enum import Enum


class FolderType(Enum):
    MOB = 1
    NPC = 2


class StanceType(Enum):
    ALL = (0, "")
    STAND = (1, "stand")
    HIT = (2, "hit1")
    JUMP = (3, "jump")
    MOVE = (4, "move")
    FLY = (MOVE, "fly")
    ATTACK = (5, "attack1")
    DIE = (6, "die1")
    ATTACK_BALL = (7, "attack.info.ball")
    ATTACK_HIT = (8, "attack.info.hit")
    ATTACK_EFFECT = (9, "attack.info.effect")

    def __init__(self, value, image_naming=None):
        self._value_ = image_naming
        if not image_naming:
            self.display_name = [t.display_name for t in self.__class__ if t != self]
        else:
            self.display_name = self._name_.lower()


images = Path("./images/sorted/Mob/")
collection = set()
base_url = URL("https://maplestory.io/api/")
desired_stances = [StanceType.STAND, StanceType.ATTACK, StanceType.MOVE]
failed = []
successful = []
STANCES = {}

for folder in images.glob("./*/"):
    collection.add(int(folder.name))

to_grab = []
with open("failed.txt", "r") as f:
    to_grab = json.loads(f.read())


async def pre_list():
    async with ClientSession() as session:
        for mob_id in to_grab:
            async with session.request("GET", (base_url / "GMS" / "")) as resp:
                if resp.status != 200:
                    continue
                mob_data = await resp.json()
                STANCES[mob_id] = {}
                for k, v in mob_data["framebooks"].items():
                    STANCES[mob_id][k] = v
            if len(STANCES) % 200 == 0:
                await sleep(3)

        for mob_id, stance in STANCES.items():
            mob_folder = images / f"{mob_id}"
            if not mob_folder.exists():
                mob_folder.mkdir(exist_ok=False)

            for k, v in stance.items():
                stance_folder = mob_folder / f"{k}"
                if not stance_folder.exists():
                    stance_folder.mkdir(exist_ok=False)
                for i in range(v):
                    async with session.request(
                        "GET",
                        (base_url / f"GMS/231/mob/{mob_id}/render/{k}/{i+1}"),
                    ) as resp:
                        if resp.status != 200:
                            failed.append(mob_id)
                            continue

                        file = stance_folder / f"{k}.{i+1}.png"

                        if not file.exists():
                            image_bytes = await resp.read()
                            file.touch()
                            async with aiofiles.open(file, "wb") as imgfl:
                                await imgfl.write(image_bytes)
                        successful.append(mob_id)


async def main():
    mobs = set()
    filtered_mobs = []
    async with ClientSession() as session:
        async with session.request("GET", base_url / "GMS/231/mob/") as resp:
            all_mobs = await resp.json()
            for mob in all_mobs:
                if int(mob["id"]) in collection:
                    continue
                mobs.add(int(mob["id"]))

        for mob_id in mobs:
            async with session.request(
                "GET", base_url / f"GMS/231/mob/{mob_id}/"
            ) as resp:
                mob_data = await resp.json()
                cleaned = {"mob_id": mob_id, "frames": {}}
                for k, v in mob_data.get("framebooks", {}).items():
                    if StanceType(k) in desired_stances:
                        cleaned["frames"][k] = mob_data["framebooks"][k]

                filtered_mobs.append(cleaned)

        for mob in filtered_mobs:
            mob_folder = images / f"{mob['mob_id']}"
            mob_folder.mkdir(exist_ok=False)
            for k, v in mob["frames"].items():
                stance_folder = mob_folder / f"{k}"
                stance_folder.mkdir(exist_ok=False)
                for i in range(v):
                    async with session.request(
                        "GET",
                        base_url / f"GMS/231/mob/{mob['mob_id']}/render/{k}/{i+1}",
                    ) as resp:
                        if resp.status == 200:
                            failed.append(mob["mob_id"])
                            continue
                        image_bytes = await resp.read()
                        file = stance_folder / f"{k}.{i+1}.png"
                        file.touch()
                        async with aiofiles.open(file, "wb") as imgfl:
                            await imgfl.write(image_bytes)
                        successful.append(mob["mob_id"])


try:
    run(pre_list())
finally:
    with open("stances.json", "a+") as f:
        f.write(json.dumps(STANCES))