import json
import os
import platform
import re
import tempfile
import zipfile

import pygame

from src.utils.paths import get_project_path, CACHE_DIR

# Textures with these parts will be excluded
EXCLUDE = {
    "bottom", "side", "front", "destroy", "on", "tip", "lit", "inner"
}

# Textures with these names will be copied to their variants (variants listed in `VARIANTS`)
# Tuples indicate a change of name for variants
TO_VARIANT = [
    ("nether_bricks", "nether_brick"),
    ("nether_bricks", "nether_brick"),
    ("oak_planks", "oak"),
    ("spruce_planks", "spruce"),
    ("birch_planks", "birch"),
    ("jungle_planks", "jungle"),
    ("acacia_planks", "acacia"),
    ("dark_oak_planks", "dark_oak"),
    ("mangrove_planks", "mangrove"),
    ("cherry_planks", "cherry"),
    ("bamboo_mosaic", "bamboo"),
    ("crimson_planks", "crimson"),
    ("warped_planks", "warped"),
    "stone",
    "cobblestone",
    "mossy_cobblestone",
    "smooth_stone",
    ("stone_bricks", "stone_brick"),
    ("mossy_stone_bricks", "mossy_stone_brick"),
    "granite",
    "polished_granite",
    "diorite",
    "polished_diorite",
    "andesite",
    "polished_andesite",
    "cobbled_deepslate",
    "polished_deepslate",
    ("deepslate_bricks", "deepslate_brick"),
    ("deepslate_tiles", "deepslate_tile"),
    ("bricks", "brick"),
    ("mud_bricks", "mud_brick"),
    "sandstone",
    "smooth_sandstone",
    "cut_sandstone",
    "red_sandstone",
    "smooth_red_sandstone",
    "cut_red_sandstone",
    "prismarine",
    ("prismarine_bricks", "prismarine_brick"),
    "dark_prismarine",
    ("red_nether_bricks", "red_nether_brick"),
    "blackstone",
    "polished_blackstone",
    ("polished_blackstone_bricks", "polished_blackstone_brick"),
    ("end_stone_bricks", "end_stone_brick"),
    ("purpur_block", "purpur"),
    ("quartz_block", "quartz"),
    "smooth_quartz",
    "cut_copper",
    "exposed_cut_copper",
    "weathered_cut_copper",
    "oxidized_cut_copper",
    "waxed_cut_copper",
    "waxed_exposed_cut_copper",
    "waxed_weathered_cut_copper",
    "waxed_oxidized_cut_copper",
]

# Variants of the textures in `TO_VARIANT`
VARIANTS = [
    "slab", "stairs",
    "wall", "fence", "fence_gate",
    "pressure_plate", "button",
    "sign", "wall_sign", "hanging_sign", "wall_hanging_sign"
]

# Colors to copy
TO_COPY = [
    ("minecraft:furnace", "minecraft:dropper"),
    ("minecraft:furnace", "minecraft:dispenser"),
    ("minecraft:furnace", "minecraft:piston"),
    ("minecraft:furnace", "minecraft:sticky_piston"),
    ("minecraft:oak_planks", "minecraft:piston_head"),
    ("minecraft:oak_planks", "minecraft:sticky_piston_head"),
    ("minecraft:torch", "minecraft:wall_torch"),
    ("minecraft:soul_torch", "minecraft:soul_wall_torch"),
    ("minecraft:redstone_torch", "minecraft:redstone_wall_torch"),
    ("minecraft:snow", "minecraft:snow_block"),
    ("minecraft:water", "minecraft:bubble_column"),
    ("minecraft:sandstone", "minecraft:smooth_sandstone"),
    ("minecraft:red_sandstone", "minecraft:smooth_red_sandstone"),
    ("minecraft:quartz_block", "minecraft:smooth_quartz"),
    ("minecraft:dripstone_block", "minecraft:pointed_dripstone"),
    ("minecraft:oak_log", "minecraft:oak_wood"),
    ("minecraft:spruce_log", "minecraft:spruce_wood"),
    ("minecraft:birch_log", "minecraft:birch_wood"),
    ("minecraft:acacia_log", "minecraft:acacia_wood"),
    ("minecraft:jungle_log", "minecraft:jungle_wood"),
    ("minecraft:cherry_log", "minecraft:cherry_wood"),
    ("minecraft:mangrove_log", "minecraft:mangrove_wood"),
    ("minecraft:dark_oak_log", "minecraft:dark_oak_wood"),
    ("minecraft:stripped_oak_log", "minecraft:stripped_oak_wood"),
    ("minecraft:stripped_spruce_log", "minecraft:stripped_spruce_wood"),
    ("minecraft:stripped_birch_log", "minecraft:stripped_birch_wood"),
    ("minecraft:stripped_acacia_log", "minecraft:stripped_acacia_wood"),
    ("minecraft:stripped_jungle_log", "minecraft:stripped_jungle_wood"),
    ("minecraft:stripped_cherry_log", "minecraft:stripped_cherry_wood"),
    ("minecraft:stripped_mangrove_log", "minecraft:stripped_mangrove_wood"),
    ("minecraft:stripped_dark_oak_log", "minecraft:stripped_dark_oak_wood"),
    ("minecraft:magma", "minecraft:magma_block"),
    ("minecraft:cut_copper", "minecraft:waxed_cut_copper"),
    ("minecraft:exposed_cut_copper", "minecraft:waxed_exposed_cut_copper"),
    ("minecraft:weathered_cut_copper", "minecraft:waxed_weathered_cut_copper"),
    ("minecraft:oxidized_cut_copper", "minecraft:waxed_oxidized_cut_copper"),
    ("minecraft:iron_block", "minecraft:heavy_weighted_pressure_plate"),
    ("minecraft:gold_block", "minecraft:light_weighted_pressure_plate"),
    ("minecraft:bricks", "minecraft:flower_pot"),
    ("minecraft:oak_log", "minecraft:campfire"),
    ("minecraft:oak_log", "minecraft:soul_campfire"),
    ("minecraft:moss_block", "minecraft:moss_carpet"),
    ("minecraft:stone", "minecraft:infested_stone"),
    ("minecraft:cobblestone", "minecraft:infested_cobblestone"),
    ("minecraft:stone_bricks", "minecraft:infested_stone_bricks"),
    ("minecraft:mossy_stone_bricks", "minecraft:infested_mossy_stone_bricks"),
    ("minecraft:chiseled_stone_bricks", "minecraft:infested_chiseled_stone_bricks"),
    ("minecraft:deepslate", "minecraft:infested_deepslate"),
    ("minecraft:infested_stone_bricks", "minecraft:cracked_infested_stone_bricks"),
    ("minecraft:cauldron", "minecraft:water_cauldron"),
    ("minecraft:cauldron", "minecraft:lava_cauldron"),
    ("minecraft:cauldron", "minecraft:powder_snow_cauldron"),
]

# Wool colors
WOOLS = [
    "red", "blue", "cyan", "gray",
    "lime", "pink", "black", "brown",
    "green", "white", "orange", "purple",
    "yellow", "magenta", "light_blue", "light_gray"
]

# Wool variants
WOOL_VARIANTS = [
    "carpet", "bed", "banner", "wall_banner"
]

# These will be removed from the textures' names
TO_STRIP = ["_top", "_stalk", "_end", "_round", "_still"]

# Minecraft version
MC_VERSION = "1.20.1"

# Minecraft root directory (platform dependent)
MC_ROOT = os.path.expanduser({
    "Linux": r"~/.minecraft",
    "Darwin": r"~/Library/Application Support/minecraft",
    "Windows": r"%APPDATA%\.minecraft"
}[platform.system()])


def extract_textures(jar_path: str, outdir: str) -> None:
    with zipfile.ZipFile(jar_path) as f:
        for info in f.infolist():
            path = info.filename
            if not re.match(r"^assets/minecraft/textures/block/[^/]+\.png$", path):
                continue
            info.filename = os.path.basename(path)
            f.extract(info, outdir)


def main() -> None:
    print(f"[1/5] Extracting Minecraft {MC_VERSION} textures")
    jar_path = os.path.join(MC_ROOT, "versions", MC_VERSION, f"{MC_VERSION}.jar")

    if not os.path.exists(jar_path):
        print(f"Couldn't find Minecraft {MC_VERSION} JAR file")
        print(f"Not at {jar_path}")
        return None

    workdir = tempfile.TemporaryDirectory()
    extract_textures(jar_path, workdir.name)

    pygame.init()
    pygame.display.set_mode((1, 1), pygame.NOFRAME | pygame.HIDDEN)

    colors = {}
    paths = os.listdir(workdir.name)
    total = len(paths)
    skipped = 0
    print("[2/5] Averaging textures")
    for i, filename in enumerate(paths):
        print(f"\r{i+1}/{total} ({i/total*100:.2f}%) {filename}", end="")
        block, ext = filename.rsplit(".", 1)
        parts = set(block.split("_"))
        if not parts.isdisjoint(EXCLUDE):
            skipped += 1
            continue

        for s in TO_STRIP:
            block = block.replace(s, "")
        img = pygame.image.load(os.path.join(workdir.name, filename)).convert_alpha()
        color = pygame.transform.average_color(img, consider_alpha=True)
        color = color[:3]
        colors[f"minecraft:{block}"] = color

    print(f"\r{total}/{total} (100%) Finished")
    print(f"Skipped {skipped} files")

    print("[3/5] Applying overrides")
    with open(get_project_path("res", "color_overrides.json"), "r") as f:
        overrides = json.load(f)
    colors.update(overrides)

    print("[4/5] Generating variants")
    for to_variant in TO_VARIANT:
        src = to_variant[0] if isinstance(to_variant, tuple) else to_variant
        dst = to_variant[1] if isinstance(to_variant, tuple) else to_variant

        for variant in VARIANTS:
            TO_COPY.append((f"minecraft:{src}", f"minecraft:{dst}_{variant}"))

    for color in WOOLS:
        for variant in WOOL_VARIANTS:
            TO_COPY.append((f"minecraft:{color}_wool", f"minecraft:{color}_{variant}"))

    for src, dst in TO_COPY:
        colors[dst] = colors[src]

    print("[5/5] Exporting colors")
    outpath = os.path.join(CACHE_DIR, "colors.json")
    with open(outpath, "w") as f:
        json.dump(colors, f, indent=4)


if __name__ == '__main__':
    main()
