#!/usr/bin/env python3
"""Generate MIAPI material JSON from the source mods' own textures.

Every material's colour palette is sampled out of the item texture shipped by
the mod that adds it, so a material always looks like the thing it is made of.

    python3 tools/generate_materials.py --jars ~/mods

The jar directory needs one jar per source mod (Mekanism, Create, Create: The
Air War, Create: Cosmonautics, Ice and Fire). Generated files are committed, so
this only needs re-running when a material definition or a source texture
changes.

Every item id in materials.py is checked against the jars; a typo fails the run
rather than shipping a material that silently refuses to load.
"""

import argparse
import json
import re
import shutil
import sys
import zipfile

from pathlib import Path

import numpy as np
from PIL import Image

sys.path.insert(0, str(Path(__file__).parent))
from materials import EXTENSIONS, MATERIALS, PACKS  # noqa: E402

NAMESPACE = "cmmodular"
PACK_FORMAT = 48  # 1.21.1

# Brightness levels of MIAPI's grayscale templates. Each gets a colour; values
# in between are interpolated by the game.
STOPS = [24, 68, 107, 150, 190, 216, 255]
# Where to sample the source texture for each stop, as a position in the
# texture's own brightness order. Tuned against MIAPI's hand-authored copper
# palette, which these reproduce to within ~2% per channel.
SAMPLE_POINTS = [0.02, 0.02, 0.20, 0.40, 0.66, 0.85, 0.98]
# The lowest stop sits below anything in the texture: item art rarely contains
# the deep shadow a full 3D model needs, and MIAPI's own palettes extrapolate
# it the same way (copper's darkest shade x0.52 is copper's darkest stop).
SHADOW_EXTENSION = 0.52
# Luminance below which a pixel is treated as outline rather than material.
OUTLINE_CUTOFF = 16

LUMA = np.array([0.2126, 0.7152, 0.0722])


class Jars:
    """Open mod jars, indexed by the asset namespaces they contain."""

    def __init__(self, folder: Path):
        self.zips = []
        self.by_namespace = {}
        for path in sorted(folder.glob("*.jar")):
            zf = zipfile.ZipFile(path)
            names = zf.namelist()
            self.zips.append((path.name, zf, set(names)))
            for name in names:
                m = re.match(r"assets/([^/]+)/", name)
                if m:
                    self.by_namespace.setdefault(m.group(1), (path.name, zf, set(names)))

    def _for(self, namespace):
        hit = self.by_namespace.get(namespace)
        if hit is None:
            raise LookupError(
                f"no jar in --jars provides the '{namespace}' namespace; "
                f"found: {', '.join(sorted(self.by_namespace)) or 'nothing'}"
            )
        return hit

    def has_item(self, item_id: str) -> bool:
        namespace, path = item_id.split(":", 1)
        try:
            _, _, names = self._for(namespace)
        except LookupError:
            return False
        return f"assets/{namespace}/models/item/{path}.json" in names

    def read_json(self, namespace, path):
        _, zf, _ = self._for(namespace)
        return json.loads(zf.read(f"assets/{namespace}/{path}"))

    def texture_of(self, item_id: str) -> str:
        """Resolve an item id to the texture id its model draws."""
        namespace, path = item_id.split(":", 1)
        _, _, names = self._for(namespace)
        seen = set()
        model_ref = f"{namespace}:item/{path}"
        while model_ref and model_ref not in seen:
            seen.add(model_ref)
            ns, mpath = model_ref.split(":", 1) if ":" in model_ref else ("minecraft", model_ref)
            key = f"assets/{ns}/models/{mpath}.json"
            if key not in names:
                break
            model = json.loads(self._for(ns)[1].read(key))
            textures = model.get("textures", {})
            for layer in ("layer0", "all", "texture", "particle"):
                if layer in textures and not textures[layer].startswith("#"):
                    return textures[layer]
            model_ref = model.get("parent")
        guess = f"{namespace}:item/{path}"
        if f"assets/{namespace}/textures/item/{path}.png" in names:
            return guess
        raise LookupError(f"could not resolve a texture for {item_id}")

    def image(self, texture_id: str) -> Image.Image:
        ns, path = texture_id.split(":", 1) if ":" in texture_id else ("minecraft", texture_id)
        _, zf, _ = self._for(ns)
        with zf.open(f"assets/{ns}/textures/{path}.png") as fh:
            img = Image.open(fh).convert("RGBA")
            img.load()
        # Animated textures stack frames vertically; the first frame is enough.
        if img.height > img.width:
            img = img.crop((0, 0, img.width, img.width))
        return img


def palette_from_image(img: Image.Image):
    """Build a MIAPI grayscale_map from a texture.

    A stop's key is the brightness of MIAPI's grayscale template; its value is
    simply what the material looks like at that point in its own shading. The
    colours are therefore taken from the texture as they are, not re-lit to
    match the stop number - which is why MIAPI's own netherite tops out at a
    murky 847a84 while iron runs to white.

    Samples are taken at percentiles of the texture's brightness order, so a
    texture with four shades still yields a full seven-stop ramp and any hue
    drift between shadow and highlight survives.
    """
    px = np.asarray(img, dtype=np.float32).reshape(-1, 4)
    opaque = px[px[:, 3] > 128][:, :3]
    if len(opaque) == 0:
        raise ValueError("texture is fully transparent")

    # Item art is usually drawn on a pure black outline. That outline is not a
    # shade of the material, and letting it win the dark stops turns the whole
    # shadow end colourless, so it is dropped - unless the material really is
    # near-black, in which case there is nothing else to sample.
    lit = opaque[opaque @ LUMA >= OUTLINE_CUTOFF]
    if len(lit) >= len(opaque) * 0.25:
        opaque = lit

    ranked = opaque[np.argsort(opaque @ LUMA)]
    window = max(1, len(ranked) // 12)

    samples = []
    for point in SAMPLE_POINTS:
        centre = int(round(point * (len(ranked) - 1)))
        lo = max(0, min(centre - window // 2, len(ranked) - window))
        # Median over a window rather than one pixel, so a stray outline pixel
        # cannot decide a whole stop.
        samples.append(np.median(ranked[lo:lo + window], axis=0))
    samples[0] = samples[1] * SHADOW_EXTENSION

    return {str(stop): "%02x%02x%02x" % tuple(int(round(c)) for c in sample)
            for stop, sample in zip(STOPS, samples)}


def build_material(spec, jars: Jars):
    texture = jars.texture_of(spec["palette_from"])
    palette = palette_from_image(jars.image(texture))

    out = {
        "translation": spec["translation"] + " ",  # trailing space: "Steel Sword"
        "groups": spec["groups"],
        "icon": {"type": "item", "item": spec["icon"]},
    }
    out.update(spec["stats"])
    out["color"] = palette["216"].upper()
    out["color_palette"] = {
        "type": "grayscale_map",
        "colors": palette,
        "filler": "interpolate",
    }
    out["textures"] = spec["textures"]
    if spec["properties"]:
        out["properties"] = spec["properties"]
    out["items"] = spec["items"]
    return out


def check_items(spec, jars: Jars):
    """Every item a material names has to exist in the jar that provides it."""
    problems = []
    for entry in spec["items"]:
        if "item" not in entry:
            continue
        if not jars.has_item(entry["item"]):
            problems.append(f"{spec['name']}: no such item {entry['item']}")
    if not jars.has_item(spec["icon"]):
        problems.append(f"{spec['name']}: icon item {spec['icon']} does not exist")
    return problems


def main():
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--jars", type=Path, required=True,
                        help="folder holding the source mods' jars")
    parser.add_argument("--out", type=Path,
                        default=Path(__file__).parent.parent / "src/main/resources/packs",
                        help="where the sub-packs are written")
    args = parser.parse_args()

    jars = Jars(args.jars)
    print(f"namespaces found: {', '.join(sorted(jars.by_namespace))}\n")

    problems = []
    for spec in MATERIALS:
        problems += check_items(spec, jars)
    if problems:
        print("aborting, item ids do not check out:", file=sys.stderr)
        for p in problems:
            print("  " + p, file=sys.stderr)
        return 1

    # Wipe first. A material that moves between packs would otherwise leave a
    # stale copy behind, and two files claiming the same material id is a silent
    # mess to debug.
    if args.out.exists():
        shutil.rmtree(args.out)

    written = 0
    for pack_id, pack in PACKS.items():
        root = args.out / pack_id
        mcmeta = {
            "pack": {
                "description": f"{pack['name']} for Truly Modular",
                "pack_format": PACK_FORMAT,
            }
        }
        root.mkdir(parents=True, exist_ok=True)
        (root / "pack.mcmeta").write_text(json.dumps(mcmeta, indent=2) + "\n")

    for spec in MATERIALS:
        material = build_material(spec, jars)
        path = (args.out / spec["pack"] / "data" / NAMESPACE / "miapi" / "materials"
                / spec["group"] / f"{spec['name']}.json")
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(material, indent=2) + "\n")
        written += 1
        swatch = " ".join(material["color_palette"]["colors"][str(s)] for s in STOPS)
        print(f"  {spec['pack']:<13} {NAMESPACE}:{spec['group']}/{spec['name']:<28} {swatch}")

    for ext in EXTENSIONS:
        path = (args.out / ext["pack"] / "data" / NAMESPACE / "miapi" / "materials"
                / f"{ext['path']}.json")
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(ext["data"], indent=2) + "\n")
        written += 1
        print(f"  {ext['pack']:<13} extends {ext['data']['key']}")

    print(f"\n{written} files written to {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())