#!/usr/bin/env python3
"""CC0 regression fixture: a skinned glTF body with an external .bin, three
external PNG textures and three materials (factor only, texture, texture x
factor), one primitive per material. Three quads side by side facing +Z.
KHR_materials_specular: Green has specularFactor 0.5; YellowTimesFactor has
specularFactor 1 with an RGBA specular texture whose alpha is 102 (0.4).

Writes tests/fixtures/bodies/multi-material/{body.gltf,body.bin,textures/*.png}.
"""
import json, os, struct, zlib

OUT = os.path.join(os.path.dirname(__file__), "..", "tests", "fixtures", "bodies", "multi-material")


def png(path, w, h, pixel, color_type=2):
    raw = b"".join(b"\0" + b"".join(bytes(pixel(x, y)) for x in range(w)) for y in range(h))
    def chunk(t, d):
        return struct.pack(">I", len(d)) + t + d + struct.pack(">I", zlib.crc32(t + d) & 0xFFFFFFFF)
    data = b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, color_type, 0, 0, 0))
    data += chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b"")
    with open(path, "wb") as f:
        f.write(data)


def main():
    os.makedirs(os.path.join(OUT, "textures"), exist_ok=True)
    # Texture A: solid green. Texture B: solid yellow (multiplied by factor 0.5,1,1).
    png(os.path.join(OUT, "textures", "green_baseColor.png"), 4, 4, lambda x, y: (0, 200, 0))
    png(os.path.join(OUT, "textures", "yellow_baseColor.png"), 4, 4, lambda x, y: (200, 200, 0))
    png(os.path.join(OUT, "textures", "specular.png"), 4, 4, lambda x, y: (0, 0, 0, 102), color_type=6)
    blob = b""
    views, accessors, prims = [], [], []

    def add(data, target=None):
        nonlocal blob
        while len(blob) % 4:
            blob += b"\0"
        view = {"buffer": 0, "byteOffset": len(blob), "byteLength": len(data)}
        if target:
            view["target"] = target
        views.append(view)
        blob += data
        return len(views) - 1

    def accessor(view, ctype, count, typ, **kw):
        accessors.append({"bufferView": view, "componentType": ctype, "count": count, "type": typ, **kw})
        return len(accessors) - 1

    for i, x0 in enumerate([-0.75, -0.25, 0.25]):
        x1 = x0 + 0.5
        pos = [(x0, 0.0, 0.0), (x1, 0.0, 0.0), (x1, 1.6, 0.0), (x0, 1.6, 0.0)]
        p = accessor(add(b"".join(struct.pack("<3f", *v) for v in pos), 34962), 5126, 4, "VEC3",
                     min=[x0, 0.0, 0.0], max=[x1, 1.6, 0.0])
        n = accessor(add(struct.pack("<3f", 0, 0, 1) * 4, 34962), 5126, 4, "VEC3")
        uv = accessor(add(b"".join(struct.pack("<2f", *t) for t in [(0, 1), (1, 1), (1, 0), (0, 0)]), 34962), 5126, 4, "VEC2")
        j = accessor(add(struct.pack("<4H", 0, 0, 0, 0) * 4, 34962), 5123, 4, "VEC4")
        w = accessor(add(struct.pack("<4f", 1, 0, 0, 0) * 4, 34962), 5126, 4, "VEC4")
        idx = accessor(add(struct.pack("<6H", 0, 1, 2, 0, 2, 3), 34963), 5123, 6, "SCALAR")
        prims.append({"attributes": {"POSITION": p, "NORMAL": n, "TEXCOORD_0": uv, "JOINTS_0": j, "WEIGHTS_0": w},
                      "indices": idx, "material": i})
    ibm = accessor(add(struct.pack("<16f", 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1)), 5126, 1, "MAT4")
    times = accessor(add(struct.pack("<2f", 0.0, 1.0)), 5126, 2, "SCALAR", min=[0.0], max=[1.0])
    rots = accessor(add(struct.pack("<8f", 0, 0, 0, 1, 0, 0, 0, 1)), 5126, 2, "VEC4")
    gltf = {
        "extensionsUsed": ["KHR_materials_specular"],
        "asset": {"version": "2.0", "generator": "paperdoll generate_multi_material_body.py (CC0)"},
        "extras": {"forward": "+Z"},
        "scene": 0,
        "scenes": [{"nodes": [0, 1]}],
        "nodes": [{"name": "Root"}, {"name": "Body", "mesh": 0, "skin": 0}],
        "skins": [{"joints": [0], "inverseBindMatrices": ibm}],
        "meshes": [{"name": "Body", "primitives": prims}],
        "materials": [
            {"name": "FactorOnly", "pbrMetallicRoughness": {"baseColorFactor": [0.2, 0.4, 0.9, 1.0]}},
            {"name": "Green", "pbrMetallicRoughness": {"baseColorTexture": {"index": 0}},
             "extensions": {"KHR_materials_specular": {"specularFactor": 0.5}}},
            {"name": "YellowTimesFactor", "pbrMetallicRoughness": {"baseColorTexture": {"index": 1},
                                                                    "baseColorFactor": [0.5, 1.0, 1.0, 1.0]},
             "extensions": {"KHR_materials_specular": {"specularTexture": {"index": 2}}}},
        ],
        "textures": [{"sampler": 0, "source": 0}, {"sampler": 0, "source": 1}, {"sampler": 0, "source": 2}],
        "samplers": [{"magFilter": 9728, "minFilter": 9728}],
        "images": [{"uri": "textures/green_baseColor.png"}, {"uri": "textures/yellow_baseColor.png"},
                   {"uri": "textures/specular.png"}],
        "animations": [{"name": "idle", "samplers": [{"input": times, "output": rots}],
                        "channels": [{"sampler": 0, "target": {"node": 0, "path": "rotation"}}]}],
        "buffers": [{"uri": "body.bin", "byteLength": len(blob)}],
        "bufferViews": views,
        "accessors": accessors,
    }
    with open(os.path.join(OUT, "body.bin"), "wb") as f:
        f.write(blob)
    with open(os.path.join(OUT, "body.gltf"), "w") as f:
        json.dump(gltf, f, indent=1)


if __name__ == "__main__":
    main()