#!/usr/bin/env python3
"""Intake for the Quaternius modular character parts.

The upstream modular pack is distributed as a Google Drive folder, so it cannot be
fetched reproducibly by CI. The owner downloads it once into the gitignored
`tmp/quaternius-intake/`; everything after that is scripted here:

    # 1. what is in the download?
    python3 tools/quaternius_parts.py inventory --dir tmp/quaternius-intake

    # 2. keep one mesh as a part (a single skinned mesh on the 62-joint rig)
    python3 tools/quaternius_parts.py extract \
      --file tmp/quaternius-intake/Male_Modular.glb \
      --mesh "Male_Hair_1" --part hair_1

    # 3. record the provenance of what was kept
    python3 tools/quaternius_parts.py verify \
      --dir assets/models/blender_src/quaternius/parts

`inventory` and `verify` run outside Blender (they parse the GLB container);
`extract` needs Blender, because it re-parents the mesh and rewrites the skin.

Part contract: one mesh, the standard 62-joint vertex-group names, materials named
after a contract slot (`skin`, `cloth`, `pants`, `hair`, `accessory`) or a
fixed-colour detail. `blender_assets.py` refuses a part that breaks it.
"""

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

ROOT = Path(__file__).resolve().parent.parent
PARTS = ROOT / "assets" / "models" / "blender_src" / "quaternius" / "parts"
INTAKE = ROOT / "tmp" / "quaternius-intake"

# The runtime rig. A part that does not carry these groups cannot be skinned by
# `NpcAvatar`, so it is not a part for this project.
REQUIRED_GROUPS = ("Hips", "Chest", "Head", "UpperArm.L", "UpperArm.R",
                   "UpperLeg.L", "UpperLeg.R")


def _glb_json(path: Path):
    """Return (gltf_json, binary_length) from a GLB container."""
    data = path.read_bytes()
    magic, version, _length = struct.unpack("<III", data[:12])
    if magic != 0x46546C67:
        raise ValueError(f"{path}: not a GLB")
    offset = 12
    gltf = None
    binary = 0
    while offset < len(data):
        chunk_length, chunk_type = struct.unpack("<II", data[offset:offset + 8])
        chunk = data[offset + 8:offset + 8 + chunk_length]
        if chunk_type == 0x4E4F534A:  # JSON
            gltf = json.loads(chunk)
        elif chunk_type == 0x004E4942:  # BIN
            binary = chunk_length
        offset += 8 + chunk_length + (-chunk_length % 4)
    if gltf is None:
        raise ValueError(f"{path}: no JSON chunk")
    return gltf, binary


def _node_names(gltf):
    return [n.get("name", "") for n in gltf.get("nodes", [])]


def _skin_bones(gltf):
    nodes = _node_names(gltf)
    return [nodes[j] for s in gltf.get("skins", []) for j in s.get("joints", [])]


def _mesh_node_names(gltf):
    """Node names that reference a mesh: the names a Blender import exposes."""
    names = []
    for node in gltf.get("nodes", []):
        if "mesh" in node:
            names.append(node.get("name", f"mesh{node['mesh']}"))
    return names


def inventory_file(path: Path):
    gltf, binary = _glb_json(path)
    bones = _skin_bones(gltf)
    materials = [m.get("name", "") for m in gltf.get("materials", [])]
    node_names = _mesh_node_names(gltf)
    meshes = []
    for index, mesh in enumerate(gltf.get("meshes", [])):
        tris = 0
        verts = 0
        skinned = True
        for prim in mesh.get("primitives", []):
            attributes = prim.get("attributes", {})
            if "indices" in prim:
                tris += gltf["accessors"][prim["indices"]]["count"] // 3
            if "POSITION" in attributes:
                verts += gltf["accessors"][attributes["POSITION"]]["count"]
            if "JOINTS_0" not in attributes:
                skinned = False
        # Prefer the node name (what Blender shows and what `extract --mesh` takes).
        name = node_names[index] if index < len(node_names) else mesh.get("name", "")
        meshes.append({"name": name, "triangles": tris,
                       "vertices": verts, "skinned": skinned})
    parts_ready = all(group in bones for group in REQUIRED_GROUPS)
    return {
        "file": str(path),
        "bytes": path.stat().st_size,
        "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
        "meshes": meshes,
        "bones": len(bones),
        "rig_ready": parts_ready,
        "materials": materials,
        "binary_bytes": binary,
    }


def inventory_fbx(files):
    """The pack ships FBX parts; name-level listing, detail comes from Blender."""
    families = {}
    for path in files:
        family = path.parent.name
        families.setdefault(family, []).append(path.stem)
    print(f"{'outfit family':24s} parts")
    for family in sorted(families):
        print(f"{family[:24]:24s} {', '.join(sorted(families[family]))}")
    print(f"\n{len(files)} FBX parts in {len(families)} families.")
    print("Next: `extract --file <fbx> --mesh <name> --part <part>`; the mesh name is the "
          "node name inside the FBX (check with `--list`).")
    return 0


def inventory(root: Path):
    fbx = sorted(p for p in root.rglob("*.fbx"))
    files = sorted(p for p in root.rglob("*.glb"))
    if not files and fbx:
        return inventory_fbx(fbx)
    if not files:
        print(f"no .glb or .fbx under {root}\n"
              f"  download the pack into {INTAKE.relative_to(ROOT)}/ first "
              f"(see TODO.md, 'Outfit part library')")
        return 1
    rows = []
    for path in files:
        try:
            rows.append(inventory_file(path))
        except (ValueError, KeyError) as error:
            print(f"skip {path.name}: {error}")
    print(f"{'file':44s} {'meshes':>6s} {'tris':>7s} {'bones':>5s} rig   meshes")
    for row in rows:
        tris = sum(m["triangles"] for m in row["meshes"])
        names = ", ".join(m["name"] for m in row["meshes"][:4])
        print(f"{Path(row['file']).name[:44]:44s} {len(row['meshes']):6d} {tris:7d} "
              f"{row['bones']:5d} {'yes' if row['rig_ready'] else 'NO ':5s} {names}")
    ready = [r for r in rows if r["rig_ready"]]
    print(f"\n{len(ready)}/{len(rows)} files carry the 62-joint rig and can be cut into parts.")
    print("Next: `extract --file <glb> --mesh <name> --part <name>` for each part to keep.")
    return 0


def verify(root: Path):
    files = sorted(p for p in root.glob("*.glb"))
    if not files:
        print(f"no part files in {root.relative_to(ROOT)}")
        return 1
    print(f"{'part':32s} {'sha256':64s} {'tris':>6s} bones")
    problems = 0
    for path in files:
        gltf, _binary = _glb_json(path)
        bones = _skin_bones(gltf)
        missing = [group for group in REQUIRED_GROUPS if group not in bones]
        tris = sum(gltf["accessors"][p["indices"]]["count"] // 3
                   for mesh in gltf.get("meshes", [])
                   for p in mesh.get("primitives", {}) if "indices" in p)
        digest = hashlib.sha256(path.read_bytes()).hexdigest()
        status = "ok" if not missing else "MISSING " + ",".join(missing)
        print(f"{path.stem[:32]:32s} {digest} {tris:6d} {len(bones):5d} {status}")
        problems += 1 if missing else 0
    print(f"\n{len(files)} parts, {problems} unusable. Record the hashes in CREDITS.md.")
    return 1 if problems else 0


def extract(source: Path, mesh_name: str, part: str, out_dir: Path):
    try:
        import bpy  # type: ignore
        from mathutils import Matrix  # type: ignore
    except ImportError:
        print("extract must run inside Blender:\n"
              f"  flatpak run org.blender.Blender --background --python {Path(__file__).name} "
              f"-- extract --file {source} --mesh {mesh_name} --part {part}",
              file=sys.stderr)
        return 2
    bpy.ops.wm.read_factory_settings(use_empty=True)
    if source.suffix.lower() == ".fbx":
        bpy.ops.import_scene.fbx(filepath=str(source))
    else:
        bpy.ops.import_scene.gltf(filepath=str(source))
    target = None
    for obj in bpy.context.scene.objects:
        if obj.type == "MESH" and obj.name.split(".")[0] == mesh_name.split(".")[0]:
            target = obj
            break
    if target is None:
        names = [o.name for o in bpy.context.scene.objects if o.type == "MESH"]
        print(f"mesh '{mesh_name}' not found; the file has: {', '.join(names)}")
        return 1
    # Keep the armature: a glTF mesh only carries JOINTS_0/WEIGHTS_0 when a skin
    # references a skeleton, so a parts file without one validates as bone-less and
    # cannot be skinned. The builder drops the part's armature when it joins the mesh
    # onto the base body, which already owns the skeleton.
    armour = next((o for o in bpy.context.scene.objects
                   if o.type == "ARMATURE" and o is not target), None)
    for obj in list(bpy.context.scene.objects):
        if obj is not target and obj is not armour:
            bpy.data.objects.remove(obj, do_unlink=True)
    for group in list(target.vertex_groups):
        if group.name not in REQUIRED_GROUPS + (
                "LowerArm.L", "LowerArm.R", "LowerLeg.L", "LowerLeg.R", "Wrist.L",
                "Wrist.R", "Foot.L", "Foot.R", "Neck", "Shoulder.L", "Shoulder.R",
                "Abdomen", "Torso", "Body", "Root", "PT.L", "PT.R"):
            target.vertex_groups.remove(group)
    out_dir.mkdir(parents=True, exist_ok=True)
    path = out_dir / f"{part}.glb"
    target.name = part
    target.data.name = part
    for obj in bpy.context.scene.objects:
        obj.select_set(obj is target or obj is armour)
    bpy.context.view_layer.objects.active = target
    bpy.ops.export_scene.gltf(filepath=str(path), export_format="GLB",
                              use_selection=True, export_apply=False,
                              export_animations=False, export_skins=True,
                              export_yup=True, export_materials="EXPORT")
    print(f"wrote {path.relative_to(ROOT)} ({len(target.data.vertices)} verts)")
    return 0


def _base_name(name: str) -> str:
    """Strip Blender's `.001` duplicate suffix."""
    head, sep, tail = name.rpartition(".")
    return head if sep and tail.isdigit() else name


def make_base(family_dir: Path, out: Path, yaw: float = 0.0):
    """Compose one outfit family from its FBX parts into a single skinned GLB base.

    The pack ships each outfit as four FBX files (Body/Legs/Feet/Head) that share one
    skeleton. Each file carries its own copy of that armature, so this keeps one and
    repoints every mesh at it before joining — otherwise `join` would leave three
    meshes bound to armatures that no longer exist.
    """
    try:
        import bpy  # type: ignore
        import bmesh  # type: ignore
    except ImportError:
        print("make-base must run inside Blender", file=sys.stderr)
        return 2
    parts = sorted(family_dir.glob("*.fbx"))
    if not parts:
        print(f"no FBX under {family_dir}")
        return 1
    bpy.ops.wm.read_factory_settings(use_empty=True)
    for path in parts:
        bpy.ops.import_scene.fbx(filepath=str(path))
    armatures = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"]
    meshes = [o for o in bpy.context.scene.objects if o.type == "MESH"]
    if not armatures or not meshes:
        print(f"{family_dir.name}: import produced no armature/mesh")
        return 1
    keep = armatures[0]
    for extra in armatures[1:]:
        bpy.data.objects.remove(extra, do_unlink=True)
    # FBX adds artefacts the runtime rig does not have: zero-length `_end` leaf bones
    # and a `neutral_bone` root. Godot's importer failed on the exported file with them
    # (`resource_format_binary` error), and the accepted rig is the 62 joints only.
    bpy.context.view_layer.objects.active = keep
    bpy.ops.object.mode_set(mode="EDIT")
    for bone in list(keep.data.edit_bones):
        if bone.name.endswith("_end") or bone.name == "neutral_bone":
            for child in list(bone.children):
                child.parent = bone.parent
                child.use_connect = False
            keep.data.edit_bones.remove(bone)
    bpy.ops.object.mode_set(mode="OBJECT")
    for mesh in meshes:
        mesh.parent = keep
        mesh.matrix_parent_inverse.identity()
        for modifier in list(mesh.modifiers):
            if modifier.type == "ARMATURE":
                mesh.modifiers.remove(modifier)
        modifier = mesh.modifiers.new("Armature", "ARMATURE")
        modifier.object = keep
    # The pack reuses one fabric material for jacket and trousers, so the joined base
    # would have no `pants` surface. Give the leg meshes their own material copies;
    # the slot map in blender_assets.py then maps those to `pants`.
    leg_copies = {}
    for mesh in meshes:
        if "Legs" not in mesh.name:
            continue
        for slot, material in enumerate(mesh.data.materials):
            if material is None:
                continue
            # `Suit` and `Suit.001` are the same fabric arriving from two FBX files;
            # normalise the name so every leg mesh shares one `Suit_legs` material.
            base = _base_name(material.name)
            if base not in leg_copies:
                copy = material.copy()
                copy.name = f"{base}_legs"
                leg_copies[base] = copy
            mesh.data.materials[slot] = leg_copies[base]
    for obj in bpy.context.selected_objects:
        obj.select_set(False)
    for mesh in meshes:
        mesh.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    if len(meshes) > 1:
        bpy.ops.object.join()
    body = bpy.context.object
    body.name = "npc_base"
    body.data.name = "npc_base"
    # Imported FBX parts share seams, which leaves zero-area triangles that the audit
    # rejects (and Godot warns about during LOD generation). Dissolve and drop them
    # after the join: nothing that small carries a meaningful weight.
    bm = bmesh.new()
    bm.from_mesh(body.data)
    bmesh.ops.dissolve_degenerate(bm, dist=2e-4, edges=bm.edges)
    bmesh.ops.triangulate(bm, faces=bm.faces[:], quad_method="BEAUTY",
                          ngon_method="BEAUTY")
    bad = [f for f in bm.faces if f.calc_area() < 1e-10]
    if bad:
        bmesh.ops.delete(bm, geom=bad, context="FACES_ONLY")
    bm.to_mesh(body.data)
    bm.free()

    # The FBX rig comes in at the pack's own scale; normalise to the 1.81 m the roster
    # uses so gear attach points and the pose maths stay valid.
    evaluated = body.evaluated_get(bpy.context.evaluated_depsgraph_get())
    zs = [(body.matrix_world @ v.co).z for v in evaluated.data.vertices]
    height = max(zs) - min(zs)
    scale = 1.81 / height if height > 0 else 1.0
    if abs(scale - 1.0) > 0.02:
        keep.scale = [component * scale for component in keep.scale]
        bpy.context.view_layer.update()
        print(f"  scaled {family_dir.name} by {scale:.3f} (was {height:.3f} m)")

    # Some families are authored facing the other way; the project's front is Godot
    # +Z (Blender -Y), so a 180 degree yaw is applied before the transforms are baked.
    # Bake the armature and mesh object transforms into their data. An FBX import
    # keeps a scale/rotation on the objects, and the glTF exporter then writes node
    # transforms that change the rest pose on import (Godot rendered the head at
    # several times its size before this).
    for obj in bpy.context.selected_objects:
        obj.select_set(False)
    for obj in (keep, body):
        obj.select_set(True)
    bpy.context.view_layer.objects.active = keep
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)

    # Some families are authored facing the other way; the project's front is Godot
    # +Z. Rotate the *armature object* and export that node transform — rotating the
    # mesh as well double-applies the skin, and baking the rotation into an already
    # bound mesh collapses the character.
    if abs(yaw) > 0.01:
        # Rotate in *data* space, after the object transforms are baked: turning the
        # exported armature node rotates with the skin and renders headless, and
        # rotating the object pair on top of a bound skin collapses it. Rotating the
        # vertices and the rest bones by the same matrix keeps the bind intact.
        from mathutils import Matrix  # type: ignore
        rotation = Matrix.Rotation(math.radians(yaw), 4, "Z")
        body.data.transform(rotation)
        bpy.context.view_layer.objects.active = keep
        bpy.ops.object.mode_set(mode="EDIT")
        for bone in keep.data.edit_bones:
            bone.head = rotation @ bone.head
            bone.tail = rotation @ bone.tail
        bpy.ops.object.mode_set(mode="OBJECT")
        bpy.context.view_layer.update()

    out.parent.mkdir(parents=True, exist_ok=True)
    for obj in bpy.context.scene.objects:
        obj.select_set(obj is body or obj is keep)
    bpy.context.view_layer.objects.active = body
    bpy.ops.export_scene.gltf(filepath=str(out), export_format="GLB",
                              use_selection=True, export_apply=False,
                              export_animations=False, export_skins=True,
                              export_yup=True, export_materials="EXPORT")
    print(f"wrote {out.relative_to(ROOT)} ({len(body.data.vertices)} verts, "
          f"{len(keep.data.bones)} bones)")
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(description="Quaternius modular part intake")
    sub = parser.add_subparsers(dest="cmd", required=True)
    inv = sub.add_parser("inventory", help="list the GLBs in a download")
    inv.add_argument("--dir", default=str(INTAKE))
    ver = sub.add_parser("verify", help="hash and validate the kept parts")
    ver.add_argument("--dir", default=str(PARTS))
    ext = sub.add_parser("extract", help="write one mesh as a part (needs Blender)")
    ext.add_argument("--file", required=True)
    ext.add_argument("--mesh", required=True)
    ext.add_argument("--part", required=True)
    ext.add_argument("--out", default=str(PARTS))
    base = sub.add_parser("make-base", help="compose one outfit family as a base GLB")
    base.add_argument("--family", required=True, help="folder under the pack's Separate Skeletal Meshes")
    base.add_argument("--out", required=True)
    base.add_argument("--yaw", type=float, default=0.0,
                      help="rotate the finished base about Z (degrees)")
    argv = list(argv if argv is not None else sys.argv[1:])
    if "--" in argv:  # Blender hands the script path over before its own `--`
        argv = argv[argv.index("--") + 1:]
    args = parser.parse_args(argv)
    if args.cmd == "inventory":
        return inventory(Path(args.dir))
    if args.cmd == "verify":
        return verify(Path(args.dir))
    if args.cmd == "make-base":
        return make_base(Path(args.family), Path(args.out), args.yaw)
    return extract(Path(args.file), args.mesh, args.part, Path(args.out))


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