# Lecture des associations materiau -> textures d'un FBX binaire, sans SDK (Tools/HairPipeline.ps1).
#
# Un FBX binaire est un arbre de noeuds ; Objects contient les Material, Texture (RelativeFilename) et Video, et
# Connections les liens "OP" (texture -> propriete d'un materiau : DiffuseColor, TransparentColor, OpacityMap...) ou
# "OO" (video -> texture). Ce script les relie et ecrit, pour chaque materiau, le fichier de chaque canal.
#
# Usage : python Tools/FbxMaterialMap.py <fichier.fbx> [<sortie.json>]
# Sortie : {"<materiau>": {"DiffuseColor": "Hair_01_Gray.png", "TransparentColor": "Hair_01A1.jpg", ...}, ...}

import json
import re
import struct
import sys
import zlib


def read_nodes(data, offset, end, wide):
    nodes = []
    while offset < end:
        if wide:
            end_offset, num_props, _ = struct.unpack_from("<QQQ", data, offset)
            offset += 24
        else:
            end_offset, num_props, _ = struct.unpack_from("<III", data, offset)
            offset += 12
        name_len = data[offset]
        offset += 1
        if end_offset == 0:
            break
        name = data[offset:offset + name_len].decode("ascii", "replace")
        offset += name_len
        props = []
        for _ in range(num_props):
            code = chr(data[offset])
            offset += 1
            if code in "YCIFDL":
                fmt = {"Y": "<h", "C": "<?", "I": "<i", "F": "<f", "D": "<d", "L": "<q"}[code]
                props.append(struct.unpack_from(fmt, data, offset)[0])
                offset += struct.calcsize(fmt)
            elif code in "SR":
                length = struct.unpack_from("<I", data, offset)[0]
                offset += 4
                raw = data[offset:offset + length]
                props.append(raw.decode("utf-8", "replace") if code == "S" else raw)
                offset += length
            elif code in "fdlib":
                count, encoding, length = struct.unpack_from("<III", data, offset)
                offset += 12 + length
                props.append(None)  # tableaux inutiles ici
            else:
                raise ValueError("type de propriete FBX inconnu : " + code)
        children = read_nodes(data, offset, end_offset, wide) if offset < end_offset else []
        nodes.append((name, props, children))
        offset = end_offset
    return nodes


def object_name(prop):
    # "Nom\x00\x01Classe" -> "Nom"
    return prop.split("\x00")[0] if isinstance(prop, str) else str(prop)


def material_map(path):
    data = open(path, "rb").read()
    if not data.startswith(b"Kaydara FBX Binary"):
        raise ValueError("FBX ASCII non gere : " + path)
    version = struct.unpack_from("<I", data, 23)[0]
    nodes = read_nodes(data, 27, len(data), version >= 7500)
    objects = next((c for n, p, c in nodes if n == "Objects"), [])
    connections = next((c for n, p, c in nodes if n == "Connections"), [])
    materials, textures, videos = {}, {}, {}
    for name, props, children in objects:
        if not props:
            continue
        oid = props[0]
        if name == "Material":
            materials[oid] = object_name(props[1])
        elif name in ("Texture", "Video"):
            filename = next((c[1][0] for c in children if c[0] in ("RelativeFilename", "FileName") and c[1]), "")
            target = textures if name == "Texture" else videos
            target[oid] = re.split(r"[\\/]", filename)[-1] if filename else object_name(props[1])
    result = {}
    for name, props, _ in connections:
        if name != "C" or len(props) < 3:
            continue
        if props[0] == "OP" and props[1] in textures and props[2] in materials:
            channel = props[3] if len(props) > 3 else "?"
            result.setdefault(materials[props[2]], {})[channel] = textures[props[1]]
    return result


if __name__ == "__main__":
    mapping = material_map(sys.argv[1])
    text = json.dumps(mapping, indent=1, ensure_ascii=False)
    if len(sys.argv) > 2:
        open(sys.argv[2], "w", encoding="utf-8").write(text)
    print(text)