pokemon-color-search/tools/download.py

"""
Manage the logic of downloading the pokedex and source images.
"""

import re
import json
import asyncio
from pathlib import Path
from dataclasses import dataclass, asdict
from collections import defaultdict

from aiohttp import ClientSession


JS_TO_JSON = re.compile(r"\b([a-zA-Z][a-zA-Z0-9]*?):")


# the dex from showdown assumes only strawberry alcremie, since
# that's what's in showdown, but we might as well add the rest
ALCREMIE_SWEETS = [
  "Strawberry", "Berry", "Love", "Star",
  "Clover", "Flower", "Ribbon",
]

# https://bulbapedia.bulbagarden.net/wiki/List_of_Pok%C3%A9mon_with_gender_differences
# there are some pokemon with notable gender diffs that the dex doesn't cover
# judgement calls made arbitrarily
GENDER_DIFFS = (
  "wobbuffet",
  "hippopotas", "hippowdon",
  "unfezant",
  "frillish", "jellicent",
  "pyroar",
  # meowstic, indeedee, basculegion, oinkologne are already handled in the dex
)


@dataclass
class Form:
  name: str
  traits: list[str]
  types: list[str]
  color: str


@dataclass
class Pokemon:
  num: int
  species: str
  forms: list[Form]


async def download_pokedex() -> dict:
  async with ClientSession() as session:
    async with session.get("https://play.pokemonshowdown.com/data/pokedex.js") as res:
      res.raise_for_status()
      text = await res.text("utf-8")
  # this is not json of course, but it's close
  # start by taking out the ; and definition
  cleaned = text.replace("exports.BattlePokedex = ", "").strip(";")
  # then convert the keys to strings
  converted = re.sub(JS_TO_JSON, lambda m: f'"{m.group(1)}":', cleaned)
  # and fix Type: Null, Farfetch'd, Sirfetch'd
  fixed = converted.replace(
    '""Type": Null"', '"Type: Null"').replace("\u2019", "'")
  # then, parse it
  return json.loads(fixed)


START_IN_BASE_FORM = (
  "Castform",    # castform can't start battle in weather forms
  "Cherrim",     # castform can't start battle in sunshine form
  "Aegislash",   # aegislash can't start battle in blade form
  "Wishiwashi",  # wishiwashi can't start battle in schooling form
  "Mimikyu",     # mimikyu can't start battle in busted form
  "Cramorant",   # cramorant can't start battle in gulping/gorging form
  "Eiscue",      # eiscue can't start battle in noice form
  "Morpeko",     # morpeko can't start battle in hangry form
  "Palafin",     # palafin can only start in zero form
  "Gimmighoul",  # gimmighoul roaming is only in PGO
)


def get_traits(species: str, form_info: dict) -> list[str]:
  traits = []
  if len(form_info.get("evos", [])) > 0:
    traits.append("nfe")
  kind = form_info["formeKind"].lower()
  if kind in ("mega", "mega-x", "mega-y", "primal"):
    traits.extend(("mega", "nostart"))
  if kind in ("gmax", "eternamax", "rapid-strike-gmax"):
    traits.extend(("gmax", "nostart"))
  if kind in ("alola", "galar", "hisui", "paldea"):
    traits.extend(("regional", kind))

  if species in START_IN_BASE_FORM and kind != "base":
    traits.append("nostart")

  # special cases
  if species == "Necrozma" and kind == "ultra":
    # necrozma can't start battle in ultra form
    traits.append("nostart")
  if species == "Tauros" and "paldea" in kind:
    # paldean tauros has dumb names
    traits.extend(("regional", "paldea"))
  if species == "Minior" and kind != "meteor":
    # minior can only start the battle in meteor form
    traits.append("nostart")
  if species == "Darmanitan" and "zen" in kind:
    # darmanitan cannot start in zen form
    traits.append("nostart")
    if "galar" in kind:
      # also there's a galar-zen form to handle
      traits.extend(("regional", "galar"))

  return sorted(set(traits))


def clean_dex(raw: dict) -> dict[int, Pokemon]:
  regrouped = defaultdict(list)

  for key, entry in raw.items():
    isNonstandard = entry.get("isNonstandard", None)
    baseSpecies = entry.get("baseSpecies", None)
    forme = entry.get("forme", None)
    if isNonstandard not in (None, "Past", "Unobtainable"):
      continue  # remove CAP etc.
    if baseSpecies in ("Pikachu", "Pichu") and forme is not None:
      continue  # remove pikachu spam + spiky ear pichu
    if forme is not None and "Totem" in forme:
      continue  # remove totem pokemon
    if baseSpecies == "Toxtricity" and forme == "Low-Key-Gmax":
      continue  # remove low-key-gmax since it is sort of a duplicate

    num = entry["num"]
    # non-cosmetic forms get separate entries automatically
    # but keeping the separate unown forms would be ridiculous
    if key != "unown" and len(cosmetic := entry.get("cosmeticFormes", [])) > 0:
      cosmetic.append(f'{entry["name"]}-{entry["baseForme"]}')
      if key == "alcremie":
        # oh god this thing
        cosmetic = [
          f"{cf}-{sweet}"
          for cf in cosmetic
          for sweet in ALCREMIE_SWEETS
        ]
      regrouped[num].extend({
        **entry,
        "forme": cf.replace(" ", "-"),
        "formeKind": "cosmetic",
      } for cf in cosmetic)
    elif key in GENDER_DIFFS:
      regrouped[num].append({
        **entry,
        "forme": f'{entry["name"]}-M',
        "formeKind": "cosmetic",
      })
      regrouped[num].append({
        **entry,
        "forme": f'{entry["name"]}-F',
        "formeKind": "cosmetic",
      })
    else:
      regrouped[num].append({
        **entry,
        "forme": entry["name"],
        "formeKind": entry.get("forme", "base").lower(),
      })

  return {
    i: Pokemon(
      num=i,
      species=(
        # doubles as an assertion that forms is not empty
        species := (forms := regrouped[i])[0].get("baseSpecies", forms[0]["name"])
      ),
      forms=[
        Form(
          name=f.get("forme", f["name"]),
          traits=get_traits(species, f),
          types=f["types"],
          color=f["color"],
        ) for f in forms
      ]
    ) for i in range(1, max(regrouped.keys()) + 1)
  }


async def load_pokedex(dex_file: Path, force_dex: bool) -> dict:
  if dex_file.is_file() and not force_dex:
    with open(dex_file) as infile:
      loaded = json.load(infile)
    dex = {
      int(num): Pokemon(
        num=entry["num"],
        species=entry["species"],
        forms=[Form(**f) for f in entry["forms"]],
      ) for num, entry in loaded.items()
    }
  else:
    # first download the pokedex
    raw_dex = await download_pokedex()
    # clean and reorganize it
    dex = clean_dex(raw_dex)
    # output dex for auditing and reloading
    with open(dex_file, "w") as out:
      json.dump({
        str(i): asdict(pkmn)
        for i, pkmn in dex.items()
      }, out, indent=2)

  return dex

SHOWDOWN_STRIP_FORM = (
  "burmy-plant",
  "shellos-west", "gastrodon-west",
  "deerling-spring", "sawsbuck-spring",
  "flabébé-red", "floette-red", "florges-red",
  "minior-red",
  "vivillon-meadow",
  "furfrou-natural",
  "alcremie-vanilla-cream-strawberry",
  "tatsugiri-curly",
)
# . for mime jr + mr mime + mr rime
# ' for farfetch'd + sirfetch'd, also oricorio-pa'u
# % for zygarde 10%
# : for type: null
# space mostly for tapus and paradox mons, but also some others
SHOWDOWN_REMOVE_SYMBOLS = ".'%: "
SHOWDOWN_REPLACEMENTS = [
  ("mega-", "mega"),  # charizard, mewtwo
  ("paldea-", "paldea"),  # tauros
  ("nidoran-m", "nidoranm"),
  ("nidoran-f", "nidoranf"),
  ("ho-oh", "hooh"),
  ("porygon-z", "porygonz"),
  ("-striped", "striped"),  # basculin special cases
  ("-galar-zen", "-galarzen"),  # darmanitan special cases
  ("é", "e"),  # flabebe
  ("vivillon-icy-snow", "vivillon-icysnow"),
  ("vivillon-high-plains", "vivillon-highplains"),
  ("furfrou-la-reine", "furfrou-lareine"),
  ("oricorio-pom-pom", "oricorio-pompom"),
  ("mo-o", "moo"),  # jangmo-o, hakamo-o, kommo-o
  ("dusk-mane", "duskmane"),  # necrozma
  ("dawn-wings", "dawnwings"),  # necrozma
  ("low-key", "lowkey"),  # toxtricity
  ("-swirl", "swirl"),  # alcremie
  ("-cream", "cream"),  # alcremie
  ("-strawberry", ""),  # alcremie
  ("rapid-strike", "rapidstrike"),  # Urshifu
  ("rapidstrike-gmax", "rapidstrikegmax"),  # Urshifu
  # dudunsparce (note - no visual diff on showdown as of 3/17)
  ("three-segment", "threesegment"),
  ("wo-chien", "wochien"),
  ("chien-pao", "chienpao"),
  ("ting-lu", "tinglu"),
  ("chi-yu", "chiyu"),
]


def get_showdown_urls(pkmn: Pokemon, form: Form) -> list[tuple[str, str]]:
  name = form.name.lower()
  if name in SHOWDOWN_STRIP_FORM:
    name = name.split("-", 1)[0]
  for c in SHOWDOWN_REMOVE_SYMBOLS:
    name = name.replace(c, "")
  for pat, ins in SHOWDOWN_REPLACEMENTS:
    name = name.replace(pat, ins)
  name = re.sub(r"-m$", "", name)  # gender diff forms
  # lol they just reuse the male back sprite for wobbuffet
  back_name = name if name != "wobbuffet-f" else "wobbuffet"
  return [
    (f"https://play.pokemonshowdown.com/sprites/ani/{name}.gif", "gif"),
    (f"https://play.pokemonshowdown.com/sprites/ani-back/{back_name}.gif", "gif"),
    (f"https://play.pokemonshowdown.com/sprites/gen5/{name}.png", "png"),
    (f"https://play.pokemonshowdown.com/sprites/gen5-back/{back_name}.png", "png"),
  ]


SEREBII_USE_SPECIES = (
  "Wobbuffet-M",
  "Shellos-West", "Gastrodon-West",
  "Hippopotas-M", "Hippowdon-M",
  "Unfezant-M",
  "Deerling-Spring", "Sawsbuck-Spring",
  "Frillish-M", "Jellicent-M",
  "Pyroar-M",
  "Furfrou-Natural",
  "Minior-Meteor",
  "Tatsugiri-Curly",
)
SEREBII_FORM_FIRST_LETTER = (
  "Deoxys-Attack", "Deoxys-Defense", "Deoxys-Speed",
  "Burmy-Plant", "Burmy-Sandy", "Burmy-Trash",
  "Cherrim-Sunshine",
  "Shellos-East", "Gastrodon-East",
  "Dialga-Origin", "Palkia-Origin", "Giratina-Origin",
  "Shaymin-Sky",
  "Unfezant-F",
  "Basculin-Blue-Striped", "Basculin-White-Striped",
  "Deerling-Summer", "Deerling-Autumn", "Deerling-Winter",
  "Sawsbuck-Summer", "Sawsbuck-Autumn", "Sawsbuck-Winter",
  "Frillish-F", "Jellicent-F",
  "Kyurem-Black", "Kyurem-White", "Keldeo-Resolute",
  "Greninja-Ash",
  "Pyroar-F",
  "Flabébé-Blue", "Flabébé-Orange", "Flabébé-White", "Flabébé-Yellow", "Flabébé-Red",
  "Floette-Blue", "Floette-Orange", "Floette-White", "Floette-Yellow", "Floette-Red",
  "Florges-Blue", "Florges-Orange", "Florges-White", "Florges-Yellow", "Florges-Red",
  "Meowstic-F",
  "Aegislash-Blade",
  "Hoopa-Unbound",
  "Lycanroc-Midnight", "Lycanroc-Dusk",
  "Wishiwashi-School",
  "Minior-Blue", "Minior-Green", "Minior-Indigo", "Minior-Orange", "Minior-Red", "Minior-Violet", "Minior-Yellow",
  "Mimikyu-Busted",
  "Magearna-Original",
  "Toxtricity-Low-Key",
  "Eiscue-Noice",
  "Indeedee-F",
  "Morpeko-Hangry",
  "Zacian-Crowned", "Zamazenta-Crowned",
  "Eternatus-Eternamax",
  "Zarude-Dada",
  "Calyrex-Ice", "Calyrex-Shadow",
  "Basculegion-F",
  "Ursaluna-Bloodmoon",
  "Oinkologne-F",
  "Maushold-Four",
  "Palafin-Hero",
  "Tatsugiri-Droopy", "Tatsugiri-Stretchy",
  "Dudunsparce-Three-Segment",
  "Gimmighoul-Roaming"
)
SEREBII_SPECIAL = {
  "Castform-Rainy": "r",
  "Castform-Snowy": "i",
  "Castform-Sunny": "s",
  # wormadam plant is default form
  "Wormadam-Sandy": "c",
  "Wormadam-Trash": "t",
  "Tauros-Paldea-Blaze": "b",
  "Tauros-Paldea-Aqua": "a",
  "Rotom-Heat": "h",
  "Rotom-Wash": "w",
  "Rotom-Frost": "f",
  "Rotom-Fan": "s",
  "Rotom-Mow": "m",
  "Darmanitan-Zen": "z",
  "Darmanitan-Galar-Zen": "gz",
  "Meloetta-Pirouette": "s",
  "Genesect-Douse": "w",
  "Genesect-Shock": "e",
  "Genesect-Burn": "f",
  "Genesect-Chill": "i",
  "Vivillon-Archipelago": "a",
  "Vivillon-Continental": "c",
  "Vivillon-Elegant": "e",
  "Vivillon-Fancy": "f",
  "Vivillon-Garden": "g",
  "Vivillon-High-Plains": "h",
  "Vivillon-Icy-Snow": "i",
  "Vivillon-Jungle": "j",
  "Vivillon-Marine": "ma",
  "Vivillon-Meadow": "m",
  "Vivillon-Modern": "mo",
  "Vivillon-Monsoon": "mon",
  "Vivillon-Ocean": "o",
  "Vivillon-Pokeball": "pb",
  "Vivillon-Polar": "p",
  "Vivillon-River": "r",
  "Vivillon-Sandstorm": "s",
  "Vivillon-Savanna": "sa",
  "Vivillon-Sun": "su",
  "Vivillon-Tundra": "t",
  "Furfrou-Dandy": "da",
  "Furfrou-Debutante": "de",
  "Furfrou-Diamond": "d",
  "Furfrou-Heart": "h",
  "Furfrou-Kabuki": "k",
  "Furfrou-La-Reine": "l",
  "Furfrou-Matron": "m",
  "Furfrou-Pharaoh": "p",
  "Furfrou-Star": "s",
  "Zygarde-10%": "10",
  "Zygarde-Complete": "c",
  "Pumpkaboo-Small": "s",  # kinda dumb but w/e
  "Pumpkaboo-Large": "l",
  "Pumpkaboo-Super": "h",
  "Gourgeist-Small": "s",
  "Gourgeist-Large": "l",
  "Gourgeist-Super": "h",
  "Oricorio-Pom-Pom": "p",
  "Oricorio-Pa'u": "pau",
  "Oricorio-Sensu": "s",
  "Necrozma-Dusk-Mane": "dm",
  "Necrozma-Dawn-Wings": "dw",
  "Necrozma-Ultra": "m",
  "Cramorant-Gulping": "gu",
  "Cramorant-Gorging": "go",
  "Sinistea-Antique": "f",
  "Polteageist-Antique": "f",
  "Urshifu-Rapid-Strike": "r",
  "Urshifu-Rapid-Strike-Gmax": "rgi",
}
SEREBII_IGNORE_MISSING = (
  "Wobbuffet-F", "Hippopotas-F", "Hippowdon-F", "Floette-Eternal",
  "Squawkabilly-Blue", "Squawkabilly-Yellow", "Squawkabilly-White",
)


def get_serebii_url(pkmn: Pokemon, form: Form) -> str | None:
  if form.name == pkmn.species or form.name in SEREBII_USE_SPECIES:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}.png"
  if form.name in SEREBII_FORM_FIRST_LETTER:
    letter = form.name.split("-", 1)[1][0].lower()
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-{letter}.png"
  if form.name in SEREBII_SPECIAL:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-{SEREBII_SPECIAL[form.name]}.png"
  if "gmax" in form.traits:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-gi.png"
  if "mega" in form.traits:
    if "Mega-X" in form.name:
      return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-mx.png"
    elif "Mega-Y" in form.name:
      return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-my.png"
    else:
      return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-m.png"
  if "alola" in form.traits:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-a.png"
  if "galar" in form.traits:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-g.png"
  if "hisui" in form.traits:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-h.png"
  if "paldea" in form.traits:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-p.png"
  if pkmn.num in (493, 773):  # arceus and silvally
    type_name = form.name.split('-')[1].lower()
    return f"https://www.serebii.net/pokemon/art/{pkmn.num}-{type_name}.png"
  if pkmn.num == 869:  # alcremie
    _, cream_name, cream_kind, candy_name = form.name.lower().split("-")
    image_name = "869-"
    if cream_name == "matcha":
      image_name += "mac"
    elif cream_name == "mint":
      image_name += "mic"
    elif cream_name != "vanilla":
      image_name += cream_name[0] + cream_kind[0]
    if candy_name != "strawberry":
      if candy_name == "love" and cream_name == "mint":
        # for some reason this name breaks the pattern
        image_name += "heart"
      else:
        image_name += candy_name
    image_name = image_name.strip("-")
    return f"https://www.serebii.net/swordshield/pokemon/{image_name}.png"
  if "-Therian" in form.name:
    return f"https://www.serebii.net/pokemon/art/{pkmn.num:03d}-s.png"
  if form.name not in SEREBII_IGNORE_MISSING:
    print(f"No Serebii URL known for {form.name}")


async def download(session: ClientSession, url: str, filename: Path, staging: bool = False) -> tuple[str, Exception | bool]:
  if filename.is_file() or staging:
    return url, False
  try:
    async with session.get(url) as res:
      res.raise_for_status()
      with open(filename, "wb") as out:
        out.write(await res.read())
  except Exception as ex:
    return url, ex
  return url, True


async def download_all_for_pokemon(pkmn: Pokemon, image_dir: Path, staging: bool = False) -> dict[str, dict[str, Exception | bool]]:
  results = defaultdict(dict)
  async with ClientSession() as session:
    for form in pkmn.forms:
      urls = []
      urls += get_showdown_urls(pkmn, form)
      urls.append((get_serebii_url(pkmn, form), "png"))
      # TODO more sources
      results[form.name].update(await asyncio.gather(*[
        download(session, url, image_dir.joinpath(
          f"{form.name}-{i}.{ext}"), staging)
        for i, (url, ext) in enumerate(urls) if url is not None
      ]))
  return results


async def download_all(image_dir: Path, pkmn: list[Pokemon], staging: bool = False) -> dict[str, dict[str, Exception | bool]]:
  image_dir.mkdir(parents=True, exist_ok=True)
  log = {}
  for p in pkmn:
    log.update(await download_all_for_pokemon(p, image_dir, staging))
  return log


KNOWN_MISSING = [
  "https://play.pokemonshowdown.com/sprites/ani/venusaur-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/venusaur-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/blastoise-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/blastoise-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/growlithe-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/growlithe-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/arcanine-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/arcanine-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/voltorb-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/voltorb-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/electrode-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/electrode-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tauros-paldeacombat.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tauros-paldeacombat.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tauros-paldeablaze.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tauros-paldeablaze.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tauros-paldeaaqua.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tauros-paldeaaqua.gif",
  "https://play.pokemonshowdown.com/sprites/ani/wooper-paldea.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/wooper-paldea.gif",
  "https://play.pokemonshowdown.com/sprites/ani/qwilfish-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/qwilfish-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/sneasel-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/sneasel-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/dialga-origin.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/dialga-origin.gif",
  "https://play.pokemonshowdown.com/sprites/ani/palkia-origin.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/palkia-origin.gif",
  "https://play.pokemonshowdown.com/sprites/ani/lilligant-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/lilligant-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/basculin-whitestriped.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/basculin-whitestriped.gif",
  "https://play.pokemonshowdown.com/sprites/ani/braviary-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/braviary-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/sliggoo-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/sliggoo-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/goodra-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/goodra-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/avalugg-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/avalugg-hisui.gif",
  "https://play.pokemonshowdown.com/sprites/ani/rillaboom-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/rillaboom-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/cinderace-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/cinderace-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/inteleon-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/urshifu-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/urshifu-gmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/urshifu-rapidstrikegmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/urshifu-rapidstrikegmax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/wyrdeer.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/wyrdeer.gif",
  "https://play.pokemonshowdown.com/sprites/ani/kleavor.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/kleavor.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ursaluna.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ursaluna.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ursaluna-bloodmoon.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ursaluna-bloodmoon.gif",
  "https://play.pokemonshowdown.com/sprites/ani/sneasler.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/sneasler.gif",
  "https://play.pokemonshowdown.com/sprites/ani/overqwil.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/overqwil.gif",
  "https://play.pokemonshowdown.com/sprites/ani/enamorus.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/enamorus.gif",
  "https://play.pokemonshowdown.com/sprites/ani/enamorus-therian.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/enamorus-therian.gif",
  "https://play.pokemonshowdown.com/sprites/ani/pawmi.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/pawmi.gif",
  "https://play.pokemonshowdown.com/sprites/ani/pawmo.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/pawmo.gif",
  "https://play.pokemonshowdown.com/sprites/ani/pawmot.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/pawmot.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tandemaus.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tandemaus.gif",
  "https://play.pokemonshowdown.com/sprites/ani/maushold.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/maushold.gif",
  "https://play.pokemonshowdown.com/sprites/ani/maushold-four.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/maushold-four.gif",
  "https://play.pokemonshowdown.com/sprites/ani/fidough.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/fidough.gif",
  "https://play.pokemonshowdown.com/sprites/ani/dachsbun.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/dachsbun.gif",
  "https://play.pokemonshowdown.com/sprites/ani/smoliv.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/smoliv.gif",
  "https://play.pokemonshowdown.com/sprites/ani/dolliv.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/dolliv.gif",
  "https://play.pokemonshowdown.com/sprites/ani/arboliva.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/arboliva.gif",
  "https://play.pokemonshowdown.com/sprites/ani/squawkabilly.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/squawkabilly.gif",
  "https://play.pokemonshowdown.com/sprites/ani/squawkabilly-blue.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/squawkabilly-blue.gif",
  "https://play.pokemonshowdown.com/sprites/ani/squawkabilly-yellow.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/squawkabilly-yellow.gif",
  "https://play.pokemonshowdown.com/sprites/ani/squawkabilly-white.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/squawkabilly-white.gif",
  "https://play.pokemonshowdown.com/sprites/ani/nacli.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/nacli.gif",
  "https://play.pokemonshowdown.com/sprites/ani/naclstack.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/naclstack.gif",
  "https://play.pokemonshowdown.com/sprites/ani/garganacl.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/garganacl.gif",
  "https://play.pokemonshowdown.com/sprites/ani/charcadet.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/charcadet.gif",
  "https://play.pokemonshowdown.com/sprites/ani/armarouge.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/armarouge.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ceruledge.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ceruledge.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tadbulb.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tadbulb.gif",
  "https://play.pokemonshowdown.com/sprites/ani/bellibolt.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/bellibolt.gif",
  "https://play.pokemonshowdown.com/sprites/ani/wattrel.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/wattrel.gif",
  "https://play.pokemonshowdown.com/sprites/ani/kilowattrel.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/kilowattrel.gif",
  "https://play.pokemonshowdown.com/sprites/ani/maschiff.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/maschiff.gif",
  "https://play.pokemonshowdown.com/sprites/ani/mabosstiff.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/mabosstiff.gif",
  "https://play.pokemonshowdown.com/sprites/ani/shroodle.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/shroodle.gif",
  "https://play.pokemonshowdown.com/sprites/ani/grafaiai.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/grafaiai.gif",
  "https://play.pokemonshowdown.com/sprites/ani/bramblin.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/bramblin.gif",
  "https://play.pokemonshowdown.com/sprites/ani/brambleghast.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/brambleghast.gif",
  "https://play.pokemonshowdown.com/sprites/ani/toedscool.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/toedscool.gif",
  "https://play.pokemonshowdown.com/sprites/ani/toedscruel.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/toedscruel.gif",
  "https://play.pokemonshowdown.com/sprites/ani/klawf.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/klawf.gif",
  "https://play.pokemonshowdown.com/sprites/ani/capsakid.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/capsakid.gif",
  "https://play.pokemonshowdown.com/sprites/ani/scovillain.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/scovillain.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tinkatink.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tinkatink.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tinkatuff.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tinkatuff.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tinkaton.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tinkaton.gif",
  "https://play.pokemonshowdown.com/sprites/ani/bombirdier.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/bombirdier.gif",
  "https://play.pokemonshowdown.com/sprites/ani/varoom.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/varoom.gif",
  "https://play.pokemonshowdown.com/sprites/ani/revavroom.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/revavroom.gif",
  "https://play.pokemonshowdown.com/sprites/ani/cyclizar.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/cyclizar.gif",
  "https://play.pokemonshowdown.com/sprites/ani/orthworm.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/orthworm.gif",
  "https://play.pokemonshowdown.com/sprites/ani/glimmet.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/glimmet.gif",
  "https://play.pokemonshowdown.com/sprites/ani/glimmora.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/glimmora.gif",
  "https://play.pokemonshowdown.com/sprites/ani/flamigo.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/flamigo.gif",
  "https://play.pokemonshowdown.com/sprites/ani/cetoddle.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/cetoddle.gif",
  "https://play.pokemonshowdown.com/sprites/ani/cetitan.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/cetitan.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tatsugiri-droopy.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tatsugiri-droopy.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tatsugiri-stretchy.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tatsugiri-stretchy.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tatsugiri.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tatsugiri.gif",
  "https://play.pokemonshowdown.com/sprites/ani/annihilape.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/annihilape.gif",
  "https://play.pokemonshowdown.com/sprites/ani/clodsire.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/clodsire.gif",
  "https://play.pokemonshowdown.com/sprites/ani/kingambit.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/kingambit.gif",
  "https://play.pokemonshowdown.com/sprites/ani/greattusk.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/greattusk.gif",
  "https://play.pokemonshowdown.com/sprites/ani/screamtail.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/screamtail.gif",
  "https://play.pokemonshowdown.com/sprites/ani/brutebonnet.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/brutebonnet.gif",
  "https://play.pokemonshowdown.com/sprites/ani/fluttermane.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/fluttermane.gif",
  "https://play.pokemonshowdown.com/sprites/ani/slitherwing.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/slitherwing.gif",
  "https://play.pokemonshowdown.com/sprites/ani/sandyshocks.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/sandyshocks.gif",
  "https://play.pokemonshowdown.com/sprites/ani/irontreads.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/irontreads.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironbundle.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironbundle.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironhands.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironhands.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironjugulis.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironjugulis.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironmoth.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironmoth.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironthorns.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironthorns.gif",
  "https://play.pokemonshowdown.com/sprites/ani/frigibax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/frigibax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/arctibax.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/arctibax.gif",
  "https://play.pokemonshowdown.com/sprites/ani/baxcalibur.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/baxcalibur.gif",
  "https://play.pokemonshowdown.com/sprites/ani/gimmighoul.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/gimmighoul.gif",
  "https://play.pokemonshowdown.com/sprites/ani/gimmighoul-roaming.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/gimmighoul-roaming.gif",
  "https://play.pokemonshowdown.com/sprites/ani/gholdengo.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/gholdengo.gif",
  "https://play.pokemonshowdown.com/sprites/ani/wochien.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/wochien.gif",
  "https://play.pokemonshowdown.com/sprites/ani/chienpao.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/chienpao.gif",
  "https://play.pokemonshowdown.com/sprites/ani/tinglu.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/tinglu.gif",
  "https://play.pokemonshowdown.com/sprites/ani/chiyu.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/chiyu.gif",
  "https://play.pokemonshowdown.com/sprites/ani/roaringmoon.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/roaringmoon.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironvaliant.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironvaliant.gif",
  "https://play.pokemonshowdown.com/sprites/ani/koraidon.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/koraidon.gif",
  "https://play.pokemonshowdown.com/sprites/ani/miraidon.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/miraidon.gif",
  "https://play.pokemonshowdown.com/sprites/ani/walkingwake.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/walkingwake.gif",
  "https://play.pokemonshowdown.com/sprites/ani/ironleaves.gif",
  "https://play.pokemonshowdown.com/sprites/ani-back/ironleaves.gif",
]
KNOWN_MISSING_PNGS = ("vivillon", "furfrou", "alcremie")


async def main(
  dex_file: Path, image_dir: Path, startIndex: int, endIndex: int,
  log_skipped: bool, force_dex: bool, dex_only: bool, staging: bool,
):
  dex = await load_pokedex(dex_file, force_dex)
  if dex_only:
    return
  log = await download_all(image_dir, (dex[i] for i in range(startIndex, endIndex + 1)), staging)
  new_downloads = 0
  for form, result in log.items():
    for url, info in result.items():
      if isinstance(info, Exception):
        if url not in KNOWN_MISSING and not (".png" in url and any(s in url for s in KNOWN_MISSING_PNGS)):
          print(f"{form}: FAILED {url} - {info}")
      elif not info:
        if staging:
          print(url)
        elif log_skipped:
          print(f"{form}: SKIPPED {url}")
      else:
        print(f"{form}: SUCCESS {url}")
        new_downloads += 1
  print(f"New Downloads: {new_downloads}")


if __name__ == "__main__":
  from argparse import ArgumentParser
  parser = ArgumentParser(
    prog="Image Retriever",
    description="Retrieve pokedex and images",
  )
  parser.add_argument(
    "-d", "--pokedex", default="data/pokedex.json", type=Path, help="Pokedex file"
  )
  parser.add_argument(
    "--refresh-dex", action="store_true", help="Update the pokedex"
  )
  parser.add_argument(
    "--pokedex-only", action="store_true", help="Quit before image download"
  )
  parser.add_argument(
    "-s", "--staging", action="store_true", help="Log URLs without downloading"
  )
  parser.add_argument(
    "-o", "--output", default="images", type=Path, help="Image output directory"
  )
  parser.add_argument(
    "--log-skipped", action="store_true", help="Log skipped images"
  )
  parser.add_argument(
    "bounds", type=lambda a: map(int, a.split("-")), default="1-151", nargs="?",
    help="Range of dex numbers to download, inclusive"
  )

  args = parser.parse_args()
  start, end = args.bounds

  asyncio.run(main(
    args.pokedex, args.output, start, end,
    args.log_skipped, args.refresh_dex, args.pokedex_only, args.staging
  ))