#!/usr/bin/env python3
"""
Rewrite glTF materials that hide their albedo in the emissive slot.

Anime models are very often exported "unlit": `KHR_materials_unlit`, a BLACK
`baseColorFactor`, and the actual texture wired to `emissiveTexture`. Renderers
that honour the unlit extension are supposed to use base colour and ignore
emission — so Blender reads black, never references the images at all, and
imports the model with `bpy.data.images` empty. The textures are not lost on
export; they are never loaded. Taila goes through the whole pipeline and comes
out a silhouette.

Patching Blender's node graph afterwards cannot fix this, because by then there
is nothing to patch — so the file is normalised BEFORE it is imported:
emissive becomes base colour, and the unlit flag is dropped. The game shades
these characters with its own toon material off ALBEDO anyway.

Pure stdlib, so it runs inside Blender's Python or out of it.
"""
import json
import os
import struct

_MAGIC = 0x46546C67
_JSON = 0x4E4F534A
_BIN = 0x004E4942


def _read_glb(path):
    with open(path, "rb") as f:
        magic, version, _total = struct.unpack("<III", f.read(12))
        if magic != _MAGIC:
            raise ValueError(f"not a GLB: {path}")
        doc = None
        chunks = []
        while True:
            header = f.read(8)
            if len(header) < 8:
                break
            length, ctype = struct.unpack("<II", header)
            data = f.read(length)
            if ctype == _JSON:
                doc = json.loads(data.decode("utf-8"))
            chunks.append((ctype, data))
    if doc is None:
        raise ValueError(f"GLB has no JSON chunk: {path}")
    return version, doc, chunks


def _write_glb(path, version, doc, chunks):
    out = []
    for ctype, data in chunks:
        if ctype == _JSON:
            data = json.dumps(doc, separators=(",", ":")).encode("utf-8")
            data += b" " * ((4 - len(data) % 4) % 4)
        else:
            data += b"\0" * ((4 - len(data) % 4) % 4)
        out.append((ctype, data))
    total = 12 + sum(8 + len(d) for _c, d in out)
    with open(path, "wb") as f:
        f.write(struct.pack("<III", _MAGIC, version, total))
        for ctype, data in out:
            f.write(struct.pack("<II", len(data), ctype))
            f.write(data)


def _is_black(colour):
    return colour is not None and max(colour[:3]) <= 0.001


def normalize_unlit(in_path, out_path):
    """Move emissive albedo into base colour. Returns how many materials changed.

    Only touches materials that are actually broken this way — a black base
    colour with something in emission. A material that already has a proper
    base colour texture is left exactly as it is.
    """
    version, doc, chunks = _read_glb(in_path)
    changed = 0
    for mat in doc.get("materials", []):
        pbr = mat.setdefault("pbrMetallicRoughness", {})
        base_factor = pbr.get("baseColorFactor", [1.0, 1.0, 1.0, 1.0])
        has_base_tex = "baseColorTexture" in pbr
        emissive_tex = mat.get("emissiveTexture")
        emissive_factor = mat.get("emissiveFactor", [0.0, 0.0, 0.0])
        if has_base_tex or not _is_black(base_factor):
            continue
        if emissive_tex is None and _is_black(emissive_factor):
            continue  # genuinely black material — leave it alone
        if emissive_tex is not None:
            pbr["baseColorTexture"] = emissive_tex
            mat.pop("emissiveTexture", None)
        alpha = base_factor[3] if len(base_factor) > 3 else 1.0
        pbr["baseColorFactor"] = [emissive_factor[0], emissive_factor[1],
                                  emissive_factor[2], alpha]
        mat["emissiveFactor"] = [0.0, 0.0, 0.0]
        # Unlit would tell the importer to ignore everything but base colour;
        # the game lights these with its own toon shader.
        ext = mat.get("extensions", {})
        ext.pop("KHR_materials_unlit", None)
        if ext:
            mat["extensions"] = ext
        else:
            mat.pop("extensions", None)
        changed += 1

    if changed:
        used = doc.get("extensionsUsed", [])
        still = any("KHR_materials_unlit" in m.get("extensions", {})
                    for m in doc.get("materials", []))
        if not still and "KHR_materials_unlit" in used:
            used.remove("KHR_materials_unlit")
            if used:
                doc["extensionsUsed"] = used
            else:
                doc.pop("extensionsUsed", None)
        _write_glb(out_path, version, doc, chunks)
    return changed


def prepare(in_path, work_dir):
    """Return a path safe to import: the original, or a normalised copy."""
    if os.path.splitext(in_path)[1].lower() != ".glb":
        return in_path
    candidate = os.path.join(
        work_dir, os.path.splitext(os.path.basename(in_path))[0] + ".albedo.glb")
    try:
        changed = normalize_unlit(in_path, candidate)
    except (OSError, ValueError, KeyError, IndexError) as e:
        print(f"WARNING: could not normalise materials in '{in_path}' ({e})")
        return in_path
    if changed:
        print(f"Moved emissive albedo into base colour on {changed} materials "
              "(model was exported unlit)")
        return candidate
    return in_path


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 3:
        print("Usage: python tools/gltf_fix.py <in.glb> <out.glb>")
        sys.exit(1)
    n = normalize_unlit(sys.argv[1], sys.argv[2])
    print(f"{n} materials rewritten")