#!/usr/bin/env python3
"""Generate two synthetic skinned test bodies as self-contained GLB files.

The meshes, skeleton and clips are procedural. Nothing is downloaded.
The figures are labelled test assets, not finished characters. glTF axes:
Y is up, the character faces -Z, bind rotations are identity.

Run from the repository root:

    python3 scripts/generate_test_bodies.py
    python3 scripts/generate_test_bodies.py --output /tmp/bodies
"""

from __future__ import annotations

import argparse
import json
import math
import struct
from pathlib import Path

JOINT_ORDER = [
    "Hips",
    "Spine",
    "Chest",
    "Neck",
    "Head",
    "LeftShoulder",
    "LeftUpperArm",
    "LeftLowerArm",
    "LeftHand",
    "RightShoulder",
    "RightUpperArm",
    "RightLowerArm",
    "RightHand",
    "LeftUpperLeg",
    "LeftLowerLeg",
    "LeftFoot",
    "RightUpperLeg",
    "RightLowerLeg",
    "RightFoot",
]


def proportions(sex: str) -> dict[str, float]:
    if sex == "male":
        return {
            "hips_y": 0.96,
            "shoulder_x": 0.20,
            "hip_x": 0.10,
            "chest_w": 0.36,
            "hip_w": 0.28,
            "upper_arm": 0.28,
            "lower_arm": 0.26,
            "upper_leg": 0.42,
            "lower_leg": 0.40,
        }
    if sex == "female":
        return {
            "hips_y": 0.88,
            "shoulder_x": 0.16,
            "hip_x": 0.13,
            "chest_w": 0.30,
            "hip_w": 0.32,
            "upper_arm": 0.26,
            "lower_arm": 0.24,
            "upper_leg": 0.40,
            "lower_leg": 0.38,
        }
    raise SystemExit(f"unknown sex: {sex}")


def skeleton(sex: str) -> list[tuple[str, str | None, tuple[float, float, float], tuple[float, float, float]]]:
    """Return (name, parent, local translation, box size) for one test body."""
    p = proportions(sex)
    shoulder = p["shoulder_x"]
    hip = p["hip_x"]
    # name, parent, local xyz, box size xyz. The box is centered on the joint.
    spec = [
        ("Hips", None, (0.0, p["hips_y"], 0.0), (p["hip_w"], 0.16, 0.18)),
        ("Spine", "Hips", (0.0, 0.14, 0.0), (p["hip_w"] * 0.9, 0.16, 0.16)),
        ("Chest", "Spine", (0.0, 0.20, 0.0), (p["chest_w"], 0.22, 0.18)),
        ("Neck", "Chest", (0.0, 0.16, 0.0), (0.10, 0.08, 0.10)),
        ("Head", "Neck", (0.0, 0.14, 0.0), (0.18, 0.22, 0.20)),
        ("LeftShoulder", "Chest", (-shoulder, 0.08, 0.0), (0.10, 0.08, 0.08)),
        ("LeftUpperArm", "LeftShoulder", (-0.08, 0.0, 0.0), (0.08, p["upper_arm"], 0.08)),
        ("LeftLowerArm", "LeftUpperArm", (0.0, -p["upper_arm"], 0.0), (0.07, p["lower_arm"], 0.07)),
        ("LeftHand", "LeftLowerArm", (0.0, -p["lower_arm"], 0.0), (0.08, 0.10, 0.04)),
        ("RightShoulder", "Chest", (shoulder, 0.08, 0.0), (0.10, 0.08, 0.08)),
        ("RightUpperArm", "RightShoulder", (0.08, 0.0, 0.0), (0.08, p["upper_arm"], 0.08)),
        ("RightLowerArm", "RightUpperArm", (0.0, -p["upper_arm"], 0.0), (0.07, p["lower_arm"], 0.07)),
        ("RightHand", "RightLowerArm", (0.0, -p["lower_arm"], 0.0), (0.08, 0.10, 0.04)),
        ("LeftUpperLeg", "Hips", (-hip, -0.06, 0.0), (0.12, p["upper_leg"], 0.12)),
        ("LeftLowerLeg", "LeftUpperLeg", (0.0, -p["upper_leg"], 0.0), (0.09, p["lower_leg"], 0.09)),
        ("LeftFoot", "LeftLowerLeg", (0.0, -p["lower_leg"], -0.06), (0.10, 0.06, 0.22)),
        ("RightUpperLeg", "Hips", (hip, -0.06, 0.0), (0.12, p["upper_leg"], 0.12)),
        ("RightLowerLeg", "RightUpperLeg", (0.0, -p["upper_leg"], 0.0), (0.09, p["lower_leg"], 0.09)),
        ("RightFoot", "RightLowerLeg", (0.0, -p["lower_leg"], -0.06), (0.10, 0.06, 0.22)),
    ]
    names = [item[0] for item in spec]
    if names != JOINT_ORDER:
        raise SystemExit(f"joint order drifted: {names}")
    return spec


def quat_x(degrees: float) -> tuple[float, float, float, float]:
    half = math.radians(degrees) * 0.5
    return (math.sin(half), 0.0, 0.0, math.cos(half))


def quat_y(degrees: float) -> tuple[float, float, float, float]:
    half = math.radians(degrees) * 0.5
    return (0.0, math.sin(half), 0.0, math.cos(half))


def walk_clips() -> dict[str, list[tuple[float, float, float, float]]]:
    times = [0.0, 0.25, 0.5, 0.75, 1.0]
    curves = {
        "LeftUpperLeg": [-28.0, 0.0, 28.0, 0.0, -28.0],
        "RightUpperLeg": [28.0, 0.0, -28.0, 0.0, 28.0],
        "LeftLowerLeg": [12.0, 8.0, 12.0, 45.0, 12.0],
        "RightLowerLeg": [12.0, 45.0, 12.0, 8.0, 12.0],
        "LeftUpperArm": [22.0, 0.0, -22.0, 0.0, 22.0],
        "RightUpperArm": [-22.0, 0.0, 22.0, 0.0, -22.0],
        "Chest": [4.0, 0.0, -4.0, 0.0, 4.0],
    }
    return {name: [quat_x(angle) for angle in angles] for name, angles in curves.items()} | {
        "_times": [(t, 0.0, 0.0, 0.0) for t in times]
    }


def attack_clips() -> dict[str, list[tuple[float, float, float, float]]]:
    # Right arm swings toward -Z (negative X rotation). Chest twists about Y.
    arm = [quat_x(a) for a in (0.0, -20.0, -100.0, -40.0, 0.0)]
    forearm = [quat_x(a) for a in (0.0, -10.0, -35.0, -15.0, 0.0)]
    chest = [quat_y(a) for a in (0.0, 6.0, 18.0, 8.0, 0.0)]
    spine = [quat_x(a) for a in (0.0, 4.0, 12.0, 4.0, 0.0)]
    return {
        "RightUpperArm": arm,
        "RightLowerArm": forearm,
        "Chest": chest,
        "Spine": spine,
        "_times": [(t, 0.0, 0.0, 0.0) for t in (0.0, 0.15, 0.35, 0.6, 0.8)],
    }


def world_translation(
    spec: list[tuple[str, str | None, tuple[float, float, float], tuple[float, float, float]]],
    name: str,
) -> tuple[float, float, float]:
    by_name = {item[0]: item for item in spec}
    x = y = z = 0.0
    seen = 0
    while name is not None:
        _n, parent, local, _size = by_name[name]
        x += local[0]
        y += local[1]
        z += local[2]
        name = parent
        seen += 1
        if seen > len(spec):
            raise SystemExit("cycle in skeleton")
    return (x, y, z)


def add_box(
    positions: list[float],
    indices: list[int],
    joints: list[int],
    weights: list[float],
    center: tuple[float, float, float],
    size: tuple[float, float, float],
    joint_index: int,
) -> None:
    hx, hy, hz = (size[0] / 2.0, size[1] / 2.0, size[2] / 2.0)
    cx, cy, cz = center
    base = len(positions) // 3
    corners = [
        (cx - hx, cy - hy, cz - hz),
        (cx + hx, cy - hy, cz - hz),
        (cx + hx, cy + hy, cz - hz),
        (cx - hx, cy + hy, cz - hz),
        (cx - hx, cy - hy, cz + hz),
        (cx + hx, cy - hy, cz + hz),
        (cx + hx, cy + hy, cz + hz),
        (cx - hx, cy + hy, cz + hz),
    ]
    for corner in corners:
        positions.extend(corner)
        joints.extend((joint_index, 0, 0, 0))
        weights.extend((1.0, 0.0, 0.0, 0.0))
    # Twelve triangles. The generator checks that each normal points outward.
    quads = [
        (0, 1, 2, 3),  # -Z
        (5, 4, 7, 6),  # +Z
        (4, 0, 3, 7),  # -X
        (1, 5, 6, 2),  # +X
        (3, 2, 6, 7),  # +Y
        (4, 5, 1, 0),  # -Y
    ]
    for a, b, c, d in quads:
        for tri in ((a, b, c), (a, c, d)):
            i0, i1, i2 = (base + tri[0], base + tri[1], base + tri[2])
            p0 = corners[tri[0]]
            p1 = corners[tri[1]]
            p2 = corners[tri[2]]
            ux, uy, uz = (p1[0] - p0[0], p1[1] - p0[1], p1[2] - p0[2])
            vx, vy, vz = (p2[0] - p0[0], p2[1] - p0[1], p2[2] - p0[2])
            nx = uy * vz - uz * vy
            ny = uz * vx - ux * vz
            nz = ux * vy - uy * vx
            mx = (p0[0] + p1[0] + p2[0]) / 3.0 - cx
            my = (p0[1] + p1[1] + p2[1]) / 3.0 - cy
            mz = (p0[2] + p1[2] + p2[2]) / 3.0 - cz
            if nx * mx + ny * my + nz * mz < 0.0:
                i1, i2 = i2, i1
            indices.extend((i0, i1, i2))


class Blob:
    def __init__(self) -> None:
        self.data = bytearray()
        self.views: list[dict[str, int]] = []
        self.accessors: list[dict[str, object]] = []

    def _align(self) -> None:
        while len(self.data) % 4:
            self.data.append(0)

    def add_f32(
        self,
        values: list[float],
        type_name: str,
        count: int,
        minimum: list[float] | None = None,
        maximum: list[float] | None = None,
    ) -> int:
        self._align()
        offset = len(self.data)
        self.data += struct.pack("<" + "f" * len(values), *values)
        view = len(self.views)
        self.views.append({"buffer": 0, "byteOffset": offset, "byteLength": len(values) * 4})
        accessor: dict[str, object] = {
            "bufferView": view,
            "componentType": 5126,
            "count": count,
            "type": type_name,
        }
        if minimum is not None and maximum is not None:
            accessor["min"] = minimum
            accessor["max"] = maximum
        self.accessors.append(accessor)
        return len(self.accessors) - 1

    def add_u16(self, values: list[int], type_name: str, count: int) -> int:
        self._align()
        offset = len(self.data)
        self.data += struct.pack("<" + "H" * len(values), *values)
        view = len(self.views)
        self.views.append({"buffer": 0, "byteOffset": offset, "byteLength": len(values) * 2})
        self.accessors.append(
            {
                "bufferView": view,
                "componentType": 5123,
                "count": count,
                "type": type_name,
            }
        )
        return len(self.accessors) - 1


def vec_min_max(values: list[float], width: int) -> tuple[list[float], list[float]]:
    minimum = [min(values[i::width]) for i in range(width)]
    maximum = [max(values[i::width]) for i in range(width)]
    return minimum, maximum


def build_glb(sex: str) -> bytes:
    spec = skeleton(sex)
    positions: list[float] = []
    indices: list[int] = []
    joint_ids: list[int] = []
    weights: list[float] = []
    for joint_index, (name, _parent, _local, size) in enumerate(spec):
        center = world_translation(spec, name)
        add_box(positions, indices, joint_ids, weights, center, size, joint_index)

    nodes: list[dict[str, object]] = []
    for name, _parent, local, _size in spec:
        children = [index for index, item in enumerate(spec) if item[1] == name]
        node: dict[str, object] = {
            "name": name,
            "translation": [local[0], local[1], local[2]],
        }
        if children:
            node["children"] = children
        nodes.append(node)
    mesh_index = len(nodes)
    nodes.append({"name": f"test-body-{sex}-mesh", "mesh": 0, "skin": 0})

    inverse_bind: list[float] = []
    for name, _parent, _local, _size in spec:
        x, y, z = world_translation(spec, name)
        inverse_bind.extend(
            [
                1.0, 0.0, 0.0, 0.0,
                0.0, 1.0, 0.0, 0.0,
                0.0, 0.0, 1.0, 0.0,
                -x, -y, -z, 1.0,
            ]
        )

    blob = Blob()
    index_acc = blob.add_u16(indices, "SCALAR", len(indices))
    pos_min, pos_max = vec_min_max(positions, 3)
    pos_acc = blob.add_f32(positions, "VEC3", len(positions) // 3, pos_min, pos_max)
    joint_acc = blob.add_u16(joint_ids, "VEC4", len(joint_ids) // 4)
    weight_acc = blob.add_f32(weights, "VEC4", len(weights) // 4)
    ibm_acc = blob.add_f32(inverse_bind, "MAT4", len(spec))

    def add_times(times: list[float]) -> int:
        return blob.add_f32(times, "SCALAR", len(times), [times[0]], [times[-1]])

    def add_outputs(samples: list[tuple[float, float, float, float]], width: int) -> int:
        flat: list[float] = []
        for sample in samples:
            flat.extend(sample[:width])
        type_name = "VEC4" if width == 4 else "VEC3"
        return blob.add_f32(flat, type_name, len(samples))

    animations = []
    for clip_name, dutch, samples_fn in (
        ("walk", "lopen", walk_clips),
        ("attack", "aanvallen", attack_clips),
    ):
        samples = samples_fn()
        times = [item[0] for item in samples["_times"]]
        time_acc = add_times(times)
        channels = []
        samplers = []
        for joint_name, values in samples.items():
            if joint_name == "_times":
                continue
            node = JOINT_ORDER.index(joint_name)
            output = add_outputs(values, 4)
            sampler_index = len(samplers)
            samplers.append(
                {"input": time_acc, "output": output, "interpolation": "LINEAR"}
            )
            channels.append(
                {"sampler": sampler_index, "target": {"node": node, "path": "rotation"}}
            )
        if times[0] != 0.0 or times[-1] <= times[0]:
            raise SystemExit(f"bad clip times for {clip_name}")
        if values_loop_broken(samples):
            raise SystemExit(f"{clip_name} does not return to its first sample")
        animations.append(
            {
                "name": clip_name,
                "channels": channels,
                "samplers": samplers,
                "extras": {"dutch": dutch, "testClip": True},
            }
        )

    color = [0.45, 0.55, 0.70, 1.0] if sex == "male" else [0.70, 0.52, 0.48, 1.0]
    gltf = {
        "asset": {
            "version": "2.0",
            "generator": "scripts/generate_test_bodies.py",
            "extras": {
                "testAsset": True,
                "finalCharacter": False,
                "label": "Synthetisch testlichaam, geen definitief karakter",
                "sex": sex,
                "license": "CC0-1.0",
                "generator": "scripts/generate_test_bodies.py",
                "up": "+Y",
                "forward": "-Z",
                "note": "Procedureel testlichaam. Geen motion capture en geen definitief karakter.",
            },
        },
        "scene": 0,
        "scenes": [{"name": f"test-body-{sex}", "nodes": [mesh_index, 0]}],
        "nodes": nodes,
        "skins": [
            {
                "name": f"test-body-{sex}-skin",
                "skeleton": 0,
                "joints": list(range(len(spec))),
                "inverseBindMatrices": ibm_acc,
            }
        ],
        "meshes": [
            {
                "name": f"test-body-{sex}-mesh",
                "primitives": [
                    {
                        "attributes": {
                            "POSITION": pos_acc,
                            "JOINTS_0": joint_acc,
                            "WEIGHTS_0": weight_acc,
                        },
                        "indices": index_acc,
                        "material": 0,
                    }
                ],
            }
        ],
        "materials": [
            {
                "name": f"test-body-{sex}-material",
                "doubleSided": True,
                "pbrMetallicRoughness": {
                    "baseColorFactor": color,
                    "metallicFactor": 0.0,
                    "roughnessFactor": 1.0,
                },
                "extras": {"testAsset": True},
            }
        ],
        "animations": animations,
        "buffers": [{"byteLength": len(blob.data)}],
        "bufferViews": blob.views,
        "accessors": blob.accessors,
    }
    return pack_glb(gltf, bytes(blob.data))


def values_loop_broken(samples: dict[str, list[tuple[float, float, float, float]]]) -> bool:
    for name, values in samples.items():
        if name == "_times":
            continue
        if values[0] != values[-1]:
            return True
    return False


def pack_glb(gltf: dict[str, object], binary: bytes) -> bytes:
    text = json.dumps(gltf, sort_keys=True, separators=(",", ":")).encode("utf-8")
    while (len(text) + 0) % 4:
        # Pad the JSON chunk with spaces, including the chunk header alignment.
        text += b" "
    while len(text) % 4:
        text += b" "
    bin_chunk = binary
    while len(bin_chunk) % 4:
        bin_chunk += b"\x00"
    json_header = struct.pack("<II", len(text), 0x4E4F534A)
    bin_header = struct.pack("<II", len(bin_chunk), 0x004E4942)
    total = 12 + len(json_header) + len(text) + len(bin_header) + len(bin_chunk)
    header = struct.pack("<III", 0x46546C67, 2, total)
    return header + json_header + text + bin_header + bin_chunk


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("tests/fixtures/bodies"),
        help="directory for test-body-male.glb and test-body-female.glb",
    )
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    for sex in ("male", "female"):
        payload = build_glb(sex)
        if payload[:4] != b"glTF":
            raise SystemExit("GLB magic missing")
        path = args.output / f"test-body-{sex}.glb"
        path.write_bytes(payload)
        print(f"wrote {path} ({len(payload)} bytes)")


if __name__ == "__main__":
    main()