"""Builds a non-humanoid figure of the Unity port from a static Meshy model and the rig and clips of
the figure it replaces (U5.5: the Glasballiste, a walking ballista; Meshy's auto rig only takes
humanoids). Source of truth: the untouched Meshy model, the rig file extracted from the Godot
reference (armature and actions only) and tools/meshy/<id>/figure.json.

  - the model: stray islands removed, scaled to the figure's width, standing on the ground,
    centred over its feet;
  - the old rig fitted to it by landmarks (leg spread, hip and weapon height, bolt tip, back of
    the stock, bow tips): every bone group is moved to the matching part of the new model, the
    location keys of the clips follow the overall scale;
  - rigid skinning: every vertex follows its nearest bone (a machine of hard parts);
  - textures, glow, faction palette and lane copy as for the Meshy figures (build_figure.py).

Standalone: blender --background --python tools/meshy/build_machine.py -- tools/meshy/g06/figure.json [--save out.blend]
"""
import importlib.util
import json
import sys
from pathlib import Path

import bpy
import numpy as np
from mathutils import Matrix, Vector

_spec = importlib.util.spec_from_file_location('lune_build_figure', str(Path(__file__).resolve().parent / 'build_figure.py'))
bf = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(bf)


def keep_main_islands(mesh, share):
    """Removes parts that hang free in the air (Meshy left a loose crystal tip): connected
    components (seam copies joined by position) below a share of all vertices."""
    import bmesh
    data = mesh.data
    count = len(data.vertices)
    parent = list(range(count))

    def find(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    for poly in data.polygons:
        verts = list(poly.vertices)
        for v in verts[1:]:
            union(verts[0], v)
    seen = {}
    for v in data.vertices:
        key = tuple(round(c, 5) for c in v.co)
        if key in seen:
            union(seen[key], v.index)
        else:
            seen[key] = v.index
    sizes = {}
    for i in range(count):
        root = find(i)
        sizes[root] = sizes.get(root, 0) + 1
    drop = {i for i in range(count) if sizes[find(i)] < share * count}
    bm = bmesh.new()
    bm.from_mesh(data)
    bm.verts.ensure_lookup_table()
    bmesh.ops.delete(bm, geom=[bm.verts[i] for i in sorted(drop)], context='VERTS')
    bm.to_mesh(data)
    bm.free()
    return len(drop)


def place(mesh, width):
    """Uniform scale to the figure's width; feet on the ground, centred over them."""
    data = mesh.data
    co = np.array([v.co[:] for v in data.vertices])
    scale = width / (co[:, 0].max() - co[:, 0].min())
    co *= scale
    low = co[:, 2].min()
    feet = co[co[:, 2] < low + 0.08 * (co[:, 2].max() - low)]
    offset = np.array([feet[:, 0].mean(), feet[:, 1].mean(), low])
    co -= offset
    data.vertices.foreach_set('co', co.ravel())
    data.update()
    return co, scale


def landmarks(co):
    height = co[:, 2].max()
    low = co[co[:, 2] < 0.35 * height]
    reach = np.abs(co[:, 0]).max()
    bow = co[np.abs(co[:, 0]) > 0.85 * reach]
    weapon_z = float(bow[:, 2].mean())
    top = co[co[:, 2] > weapon_z - 0.12 * height]
    return {'height': float(height), 'leg_x': float(np.abs(low[:, 0]).mean()), 'weapon_z': weapon_z,
            'tip_y': float(co[:, 1].min()), 'back_y': float(top[:, 1].max()), 'bow_x': float(reach), 'bow_y': float(bow[:, 1].mean())}


def fit_rig(arm, old, new):
    """Moves every bone group of the old rig onto the matching part of the new model."""
    s = new['weapon_z'] / old['weapon_z']
    hip_old, hip_new = old['hip_z'], new['weapon_z'] * old['hip_z'] / old['weapon_z']
    front = new['tip_y'] / old['tip_y']
    back = new['back_y'] / old['back_y']

    def weapon(p):
        y = p.y * (front if p.y < 0 else back)
        return Vector((p.x * s, y, p.z - old['weapon_z'] + new['weapon_z']))

    def legs(p):
        return Vector((p.x * new['leg_x'] / old['leg_x'], p.y * s, p.z * hip_new / hip_old))

    def spine(p):
        t = (p.z - hip_old) / (old['weapon_z'] - hip_old)
        return Vector((p.x * s, p.y * s, hip_new + t * (new['weapon_z'] - hip_new)))

    def limb(p, tip):
        q = weapon(p)
        if tip:
            q.x = np.sign(p.x) * new['bow_x']
            q.y = new['bow_y']
        else:
            q.x = p.x * new['bow_x'] / old['bow_x']
        return q

    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode='EDIT')
    for bone in arm.data.edit_bones:
        name = bone.name
        head, tail = bone.head.copy(), bone.tail.copy()
        if name.startswith(('Thigh', 'Shin', 'Tarsus', 'Foot', 'Pelvis')):
            bone.head, bone.tail = legs(head), legs(tail)
        elif name == 'Spine':
            bone.head, bone.tail = spine(head), spine(tail)
        elif name.startswith('LimbTip'):
            bone.head, bone.tail = limb(head, False), limb(tail, True)
        elif name.startswith('Limb'):
            bone.head, bone.tail = limb(head, False), limb(tail, False)
        elif name != 'Root':
            bone.head, bone.tail = weapon(head), weapon(tail)
    bpy.ops.object.mode_set(mode='OBJECT')
    return s


def fcurves(action):
    if hasattr(action, 'layers') and len(action.layers):
        for layer in action.layers:
            for strip in layer.strips:
                for bag in strip.channelbags:
                    yield from bag.fcurves
    elif hasattr(action, 'fcurves'):
        yield from action.fcurves


def scale_locations(action, s):
    for curve in fcurves(action):
        if curve.data_path.endswith('location'):
            for key in curve.keyframe_points:
                key.co.y *= s
                key.handle_left.y *= s
                key.handle_right.y *= s


def rigid_skin(mesh, arm, regions):
    """Every vertex follows its nearest deforming bone (distance to the bone segment). A bone with
    a region (figure.json "bone_regions", e.g. the bolt only in front of the bow) takes only
    vertices inside it, so a long rail bone does not carry the whole stock."""
    bones = [b for b in arm.data.bones if b.use_deform and b.name != 'Root']
    co = np.array([v.co[:] for v in mesh.data.vertices])
    best = np.full(len(co), np.inf)
    owner = np.zeros(len(co), dtype=int)
    for index, bone in enumerate(bones):
        a = np.array((arm.matrix_world @ bone.head_local)[:])
        b = np.array((arm.matrix_world @ bone.tail_local)[:])
        ab = b - a
        t = np.clip(((co - a) @ ab) / max(ab @ ab, 1e-9), 0.0, 1.0)
        d = np.linalg.norm(co - (a + t[:, None] * ab), axis=1)
        region = regions.get(bone.name, {})
        for axis, name in enumerate('xyz'):
            if name + '_max' in region:
                d[co[:, axis] > region[name + '_max']] = np.inf
            if name + '_min' in region:
                d[co[:, axis] < region[name + '_min']] = np.inf
        closer = d < best
        best[closer] = d[closer]
        owner[closer] = index
    counts = {}
    for index, bone in enumerate(bones):
        members = np.nonzero(owner == index)[0].tolist()
        group = mesh.vertex_groups.new(name=bone.name)
        if members:
            group.add(members, 1.0, 'REPLACE')
        counts[bone.name] = len(members)
    modifier = mesh.modifiers.new('Armature', 'ARMATURE')
    modifier.object = arm
    mesh.parent = arm
    mesh.matrix_parent_inverse = Matrix.Identity(4)
    return counts


def build(figure_path, source_root, unity_root):
    figure = json.loads(figure_path.read_text())
    bpy.ops.wm.read_factory_settings(use_empty=True)
    scene = bpy.context.scene
    scene.render.fps = bf.FPS
    scene.render.fps_base = 1.0
    scene.unit_settings.system = 'METRIC'
    scene.unit_settings.scale_length = 1.0

    objects = bf.import_glb(source_root / figure['model'])
    mesh = next(o for o in objects if o.type == 'MESH')
    for obj in objects:
        if obj is not mesh:
            bpy.data.objects.remove(obj, do_unlink=True)
    bpy.ops.object.select_all(action='DESELECT')
    mesh.select_set(True)
    bpy.context.view_layer.objects.active = mesh
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    dropped = keep_main_islands(mesh, figure.get('min_island_share', 0.01))
    co, model_scale = place(mesh, figure['width'])
    new = landmarks(co)

    with bpy.data.libraries.load(str(source_root / figure['rig_blend'])) as (src, dst):
        dst.objects = [figure['armature']]
        dst.actions = list(src.actions)
    arm = dst.objects[0]
    scene.collection.objects.link(arm)
    if arm.animation_data:
        arm.animation_data.action = None
    for action in list(bpy.data.actions):
        if action.name not in figure['clips']:
            bpy.data.actions.remove(action)
    s = fit_rig(arm, figure['rig_landmarks'], new)
    for action in bpy.data.actions:
        scale_locations(action, s)
        action.use_fake_user = True
    mesh.name = mesh.data.name = figure['id'].capitalize() + 'Body'
    weights = rigid_skin(mesh, arm, figure.get('bone_regions', {}))

    material = mesh.material_slots[0].material
    maps = bf.textures(material, str(unity_root / figure['textures']), figure['glow'], figure.get('palette', {}))
    material.name = figure['material']
    for node in list(material.node_tree.nodes):
        if node.type not in ('BSDF_PRINCIPLED', 'OUTPUT_MATERIAL'):
            material.node_tree.nodes.remove(node)
    for image in list(bpy.data.images):
        bpy.data.images.remove(image)
    lod = bf.lane_lod(mesh, figure['lod']) if 'lod' in figure else None
    scene.frame_set(0)
    clips = {a.name: {'seconds': round((a.frame_range[1] - a.frame_range[0]) / bf.FPS, 4), 'loop': bool(a.get('loop', False)),
                      'contact_seconds': float(a.get('contact_seconds', -1.0))} for a in bpy.data.actions}
    return {'model_scale': round(model_scale, 4), 'dropped_vertices': dropped, 'landmarks': {k: round(v, 4) for k, v in new.items()},
            'rig_scale': round(s, 4), 'weights': weights, 'clips': clips, 'textures': maps, 'lod': lod}


if __name__ in ('__main__', '__lune_generator__'):
    bf.main(build, globals())