"""Builds a skinned figure of the Unity port from Meshy downloads (U5.5 pilot, docs/port/U5-VISUAL.md).

Source of truth are the untouched Meshy results in assets-source/meshy/<id>/ (the rigged character
and the merged animation GLB) and the figure file tools/meshy/<id>/figure.json. This script turns
them into the scene the FBX exporter expects (tools/blender/lune_fbx_export.py runs it as the
generator of the manifest entry):

  - the rigged mesh and its armature, scale applied (Meshy delivers centimetres under a 0.01 root),
    stray helper objects removed, a "Root" bone at the origin above Hips (pipeline rule root_bone);
  - one action per game clip (idle, walk, attack_a, attack_b, hit, ability, death), resampled from
    a library clip: frame range, new length, loop flag and the contact time of the hit;
  - one material without image nodes, named after the figure's Surface material in Unity;
  - the textures for Unity: albedo with a glow mask in alpha (lava is derived from colour and
    metalness: Meshy 7.1 delivers no emission map), its rock and metal brought to the faction
    colours of the replaced figure ("palette"), ORM (glTF metal/rough: G rough, B metal) and the
    tangent normal map, written to Assets/Luneraria/Art/Textures/<textures>-*-2k.png.

Standalone (preview): blender --background --python tools/meshy/build_figure.py -- tools/meshy/g01/figure.json [--save out.blend]
"""
import bpy
import json
import math
import sys
from pathlib import Path

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

FPS = 24
CLIP_ORDER = ['idle', 'walk', 'attack_a', 'attack_b', 'hit', 'ability', 'death']


def options(env=None):
    """Figure file and roots: from the exporter (the generator's globals) or the command line."""
    g = env if env is not None else globals()
    if 'LUNE_ARGS' in g:
        temp_root = Path(g['__file__']).resolve().parents[2]
        return temp_root / g['LUNE_ARGS'][0], temp_root, Path(g['LUNE_UNITY_ROOT']), None
    argv = sys.argv[sys.argv.index('--') + 1:] if '--' in sys.argv else []
    root = Path(g['__file__']).resolve().parents[2]
    save = argv[argv.index('--save') + 1] if '--save' in argv else None
    return root / argv[0], root, root, save


def import_glb(path):
    before = set(bpy.data.objects)
    bpy.ops.import_scene.gltf(filepath=str(path))
    return [o for o in bpy.data.objects if o not in before]


def sample_clip(arm, action, frames, seconds):
    """Pose of every bone (matrix_basis) at the frames of the new clip, read from the source action."""
    arm.animation_data.action = action
    if len(action.slots):
        arm.animation_data.action_slot = action.slots[0]
    first, last = frames
    count = max(1, int(round(seconds * FPS)))
    scene = bpy.context.scene
    samples = []
    for t in range(count + 1):
        source = first + (last - first) * t / count
        whole = math.floor(source)
        scene.frame_set(int(whole), subframe=source - whole)
        samples.append({pb.name: pb.matrix_basis.copy() for pb in arm.pose.bones})
    return samples


def write_clip(arm, name, samples, unit):
    """New action from sampled poses; translations scaled to metres, quaternions kept continuous."""
    action = bpy.data.actions.new(name)
    # Kept when another clip is assigned (the FBX exporter bakes every action of the file).
    action.use_fake_user = True
    arm.animation_data.action = action
    previous = {}
    for frame, pose in enumerate(samples):
        for pb in arm.pose.bones:
            basis = pose.get(pb.name)
            if basis is None:
                continue
            loc, rot, scale = basis.decompose()
            if pb.name in previous and previous[pb.name].dot(rot) < 0.0:
                rot.negate()
            previous[pb.name] = rot
            pb.rotation_mode = 'QUATERNION'
            pb.location = loc * unit
            pb.rotation_quaternion = rot
            pb.scale = scale
            pb.keyframe_insert('location', frame=frame)
            pb.keyframe_insert('rotation_quaternion', frame=frame)
            pb.keyframe_insert('scale', frame=frame)
    action.use_frame_range = True
    action.frame_start = 0
    action.frame_end = len(samples) - 1
    return action


def add_root(arm, name):
    """A root bone at the origin above the delivered root (Hips); its pose stays the rest pose."""
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode='EDIT')
    bones = arm.data.edit_bones
    old_roots = [b for b in bones if b.parent is None]
    root = bones.new(name)
    root.head = (0.0, 0.0, 0.0)
    root.tail = (0.0, 0.0, 0.1)
    root.use_deform = True
    for bone in old_roots:
        bone.parent = root
        bone.use_connect = False
    bpy.ops.object.mode_set(mode='OBJECT')


def smoothstep(edge0, edge1, x):
    t = np.clip((x - edge0) / (edge1 - edge0), 0.0, 1.0)
    return t * t * (3.0 - 2.0 * t)


def pixels(image):
    width, height = image.size
    data = np.empty(width * height * 4, dtype=np.float32)
    image.pixels.foreach_get(data)
    return data.reshape(height, width, 4)


def save_png(data, path, channels):
    """8-bit PNG of Blender pixels (rows bottom up), byte-identical for identical input: the
    smaller of the PNG filters Sub and Up for all rows, zlib level 9, no metadata."""
    import struct
    import zlib
    rows = np.flipud(np.clip(np.rint(data[..., :channels] * 255.0), 0, 255).astype(np.uint8))
    height, width = rows.shape[:2]
    flat = rows.reshape(height, width * channels).astype(np.int16)
    sub = flat.copy()
    sub[:, channels:] -= flat[:, :-channels]
    up = flat.copy()
    up[1:] -= flat[:-1]
    best = None
    for kind, filtered in ((1, sub), (2, up)):
        raw = np.concatenate([np.full((height, 1), kind, np.uint8), (filtered % 256).astype(np.uint8)], axis=1)
        packed = zlib.compress(raw.tobytes(), 9)
        if best is None or len(packed) < len(best):
            best = packed

    def chunk(tag, body):
        return struct.pack('>I', len(body)) + tag + body + struct.pack('>I', zlib.crc32(tag + body) & 0xffffffff)

    header = struct.pack('>IIBBBBB', width, height, 8, 6 if channels == 4 else 2, 0, 0, 0)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(b'\x89PNG\r\n\x1a\n' + chunk(b'IHDR', header) + chunk(b'IDAT', best) + chunk(b'IEND', b''))


def glow_mask(albedo, metal, rule):
    """Lava: saturated red to yellow, bright, not metal (the gold shares the hue but is metallic)."""
    rgb = albedo[..., :3]
    value = rgb.max(axis=-1)
    low = rgb.min(axis=-1)
    saturation = np.where(value > 1e-4, (value - low) / np.maximum(value, 1e-4), 0.0)
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    hue = (np.degrees(np.arctan2(math.sqrt(3.0) * (g - b), 2.0 * r - g - b)) + 360.0) % 360.0
    hue = np.where(hue > 300.0, hue - 360.0, hue)
    low_hue, high_hue = rule['hue']
    in_hue = smoothstep(low_hue - 12.0, low_hue, hue) * (1.0 - smoothstep(high_hue, high_hue + 12.0, hue))
    mask = in_hue * smoothstep(rule['saturation'] - .1, rule['saturation'] + .1, saturation) \
        * smoothstep(rule['value'] - .1, rule['value'] + .15, value) * (1.0 - smoothstep(rule['metal_below'] - .2, rule['metal_below'], metal))
    return np.clip(mask, 0.0, 1.0)


def srgb_to_linear(c):
    return np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)


def linear_to_srgb(c):
    c = np.clip(c, 0.0, 1.0)
    return np.where(c <= 0.0031308, c * 12.92, 1.055 * np.power(c, 1.0 / 2.4) - 0.055)


def hex_linear(code):
    return srgb_to_linear(np.array([int(code[i:i + 2], 16) / 255.0 for i in (0, 2, 4)]))


def palette(albedo, metal, glow, rules):
    """The faction palette of the figure it replaces: the mean colour (linear) of the rock texels
    and of the metal texels is scaled to the colour of the Phase-4 material (basalt, gold), the
    variation of the map stays. strength blends the gain (1 = the old mean exactly). Glowing texels
    keep their colour. Returns the gains for the report."""
    rgb = srgb_to_linear(albedo[..., :3])
    metalness = smoothstep(0.3, 0.7, metal)
    weights = {'rock': (1.0 - metalness) * (1.0 - glow), 'metal': metalness * (1.0 - glow)}
    factor = np.repeat(glow[..., None], 3, axis=-1)
    gains = {}
    for name, weight in weights.items():
        rule = rules.get(name)
        if not rule:
            factor += weight[..., None]
            continue
        mean = (rgb * weight[..., None]).reshape(-1, 3).sum(0) / max(float(weight.sum()), 1e-6)
        gain = np.power(hex_linear(rule['color']) / np.maximum(mean, 1e-5), rule.get('strength', 1.0))
        factor += weight[..., None] * gain
        gains[name] = [round(float(g), 4) for g in gain]
    albedo[..., :3] = linear_to_srgb(rgb * factor)
    return gains


def textures(material, prefix, rule, colours):
    """Albedo + glow, ORM and normal map of the figure from the images of its glTF material."""
    found = {}
    for node in material.node_tree.nodes:
        if node.type != 'TEX_IMAGE' or not node.image:
            continue
        for link in node.outputs['Color'].links:
            target = link.to_node
            if target.type == 'BSDF_PRINCIPLED' and link.to_socket.name == 'Base Color':
                found['albedo'] = node.image
            elif target.type == 'SEPARATE_COLOR':
                found['orm'] = node.image
            elif target.type == 'NORMAL_MAP':
                found['normal'] = node.image
    missing = [k for k in ('albedo', 'orm', 'normal') if k not in found]
    if missing:
        raise RuntimeError('material %s lacks %s' % (material.name, missing))
    albedo = pixels(found['albedo'])
    orm = pixels(found['orm'])
    normal = pixels(found['normal'])
    mask = glow_mask(albedo, orm[..., 2], rule)
    gains = palette(albedo, orm[..., 2], mask, colours) if colours else {}
    albedo[..., 3] = mask
    orm[..., 0] = 1.0
    save_png(albedo, Path(prefix + '-albedo-2k.png'), 4)
    save_png(orm, Path(prefix + '-orm-2k.png'), 3)
    save_png(normal, Path(prefix + '-normal-2k.png'), 3)
    return {'glow_share': round(float(mask.mean()), 4), 'palette_gains': gains, 'size': list(found['albedo'].size)}


def attach_prop(arm, prop, source_root, unity_root, figure):
    """A weapon generated as its own Meshy model (the auto rig would bind a held lance to hand and
    leg at once): scaled to its length, laid into the fist of the bone in the rest pose (axis along
    the figure's forward, perpendicular to the hand bone, tip forward, grip at a share of its
    length from the butt) and skinned to that bone alone, so every clip carries it."""
    objects = import_glb(source_root / prop['glb'])
    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)
    data = mesh.data
    coords = np.empty(len(data.vertices) * 3)
    data.vertices.foreach_get('co', coords)
    coords = coords.reshape(-1, 3)
    axis = int(np.argmax(coords.max(0) - coords.min(0)))
    low, high = coords[:, axis].min(), coords[:, axis].max()
    tip_positive = prop.get('tip', '+') == '+'
    # Lance frame: u along the shaft from butt to tip, origin at the grip.
    u = np.zeros(3)
    u[axis] = 1.0 if tip_positive else -1.0
    butt = low if tip_positive else high
    scale = prop['length'] / (high - low)
    grip = butt + (high - low) * prop.get('grip', 0.4) * (1 if tip_positive else -1)
    centre = coords.mean(0)
    centre[axis] = grip
    bone = arm.data.bones[prop['bone']]
    head = arm.matrix_world @ bone.head_local
    tail = arm.matrix_world @ bone.tail_local
    d = (tail - head).normalized()
    forward = Vector((0.0, -1.0, 0.0))
    shaft = (forward - d * forward.dot(d)).normalized()
    side = d.cross(shaft).normalized()
    target = head + d * prop.get('palm', 0.07)
    # Rotation taking the lance's own axes (u, two others) to (shaft, d, side).
    a = [i for i in range(3) if i != axis]
    src = np.zeros((3, 3))
    src[:, 0] = u
    src[a[0], 1] = 1.0
    src[a[1], 2] = 1.0
    dst = np.array([shaft[:], d[:], side[:]]).T
    rot = dst @ np.linalg.inv(src)
    if np.linalg.det(rot) < 0:
        dst[:, 2] *= -1
        rot = dst @ np.linalg.inv(src)
    moved = ((coords - centre) * scale) @ rot.T + np.array(target[:])
    data.vertices.foreach_set('co', moved.ravel())
    data.update()
    mesh.name = data.name = prop['name']
    group = mesh.vertex_groups.new(name=prop['bone'])
    group.add(list(range(len(data.vertices))), 1.0, 'REPLACE')
    modifier = mesh.modifiers.new('Armature', 'ARMATURE')
    modifier.object = arm
    mesh.parent = arm
    mesh.matrix_parent_inverse = Matrix.Identity(4)
    material = mesh.material_slots[0].material
    maps = textures(material, str(unity_root / prop['textures']), figure['glow'], figure.get('palette', {}))
    material.name = prop['material']
    for node in list(material.node_tree.nodes):
        if node.type not in ('BSDF_PRINCIPLED', 'OUTPUT_MATERIAL'):
            material.node_tree.nodes.remove(node)
    return mesh, {'name': prop['name'], 'triangles': sum(len(poly.vertices) - 2 for poly in data.polygons), 'textures': maps}


def lane_lod(body, rule):
    """The lane-distance copy (LOD1): the seam copies welded (Meshy splits the mesh along its UV
    seams), decimated to the target triangle count, sharp edges from the angle instead of Meshy's
    split normals. Decimation keeps UVs, material and vertex groups (the skin weights).
    The meshes are named <body>_LOD0 / _LOD1, Unity builds the LODGroup from the names."""
    import bmesh
    lod = body.copy()
    lod.data = body.data.copy()
    body.users_collection[0].objects.link(lod)
    mesh = lod.data
    if mesh.has_custom_normals:
        bpy.context.view_layer.objects.active = lod
        with bpy.context.temp_override(object=lod, active_object=lod, selected_objects=[lod]):
            bpy.ops.mesh.customdata_custom_splitnormals_clear()
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bmesh.ops.remove_doubles(bm, verts=bm.verts, dist=rule.get('weld', 1e-4))
    bm.to_mesh(mesh)
    bm.free()
    before = sum(len(poly.vertices) - 2 for poly in mesh.polygons)
    decimate = lod.modifiers.new('LaneLod', 'DECIMATE')
    decimate.decimate_type = 'COLLAPSE'
    decimate.ratio = min(1.0, rule['triangles'] / max(before, 1))
    decimate.use_collapse_triangulate = True
    while lod.modifiers[0] != decimate:
        lod.modifiers.move(lod.modifiers.find(decimate.name), 0)
    with bpy.context.temp_override(object=lod, active_object=lod):
        bpy.ops.object.modifier_apply(modifier=decimate.name)
    # Collapsing can leave zero-area triangles; the FBX round trip drops them, so they go here ...
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bmesh.ops.dissolve_degenerate(bm, edges=bm.edges, dist=1e-6)
    bmesh.ops.triangulate(bm, faces=[f for f in bm.faces if len(f.verts) > 3])
    # ... and faces folded onto the same three vertices (the FBX import keeps only one).
    seen = set()
    twins = []
    for face in bm.faces:
        key = tuple(sorted(v.index for v in face.verts))
        if key in seen:
            twins.append(face)
        seen.add(key)
    bmesh.ops.delete(bm, geom=twins, context='FACES_ONLY')
    bm.to_mesh(mesh)
    bm.free()
    for poly in mesh.polygons:
        poly.use_smooth = True
    mesh.set_sharp_from_angle(angle=math.radians(rule.get('sharp_degrees', 50)))
    name = body.name
    body.name = body.data.name = name + '_LOD0'
    lod.name = lod.data.name = name + '_LOD1'
    after = sum(len(poly.vertices) - 2 for poly in mesh.polygons)
    return {'lod0_triangles': sum(len(poly.vertices) - 2 for poly in body.data.polygons),
            'lod1_triangles': after, 'lod1_vertices': len(mesh.vertices), 'welded_triangles': before,
            'lod0_screen': rule.get('lod0_screen', 0.1)}


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 = FPS
    scene.render.fps_base = 1.0
    scene.unit_settings.system = 'METRIC'
    scene.unit_settings.scale_length = 1.0

    # Clips: sampled from the animation file before its objects go (they share the bone names).
    loose = import_glb(source_root / figure['animations'])
    arm_anim = next(o for o in loose if o.type == 'ARMATURE')
    unit = arm_anim.matrix_world.to_scale().x
    sampled = {}
    for name in CLIP_ORDER:
        spec = figure['clips'][name]
        source = bpy.data.actions[spec['source']]
        first, last = spec.get('frames', list(source.frame_range))
        seconds = spec.get('seconds', (last - first) / FPS)
        samples = sample_clip(arm_anim, source, (first, last), seconds)
        if spec.get('in_place'):
            # The lunge of a library clip without its step: Hips keeps its horizontal place
            # (bone X and Z; the bone points up), so the figure strikes from where it stands.
            start = samples[0]['Hips'].to_translation()
            for pose in samples:
                hips = pose['Hips']
                hips[0][3] = start.x
                hips[2][3] = start.z
        calm = figure.get('calm', {})
        if name in calm.get('clips', []):
            # Upright stance (U5.5 review): the library idle looks around and the orc walk hunches
            # 30-50 degrees forward, which hides the cauldrons, ovens and shields of figures without
            # a real head. Listed bones keep only a share of their rotation (0 = rest pose).
            for pose in samples:
                for bone, share in calm['bones'].items():
                    basis = pose.get(bone)
                    if basis is None:
                        continue
                    loc, rot, scale = basis.decompose()
                    rot = Quaternion().slerp(rot, share)
                    pose[bone] = Matrix.LocRotScale(loc, rot, scale)
        sampled[name] = (spec, samples, (first, last))
    for obj in loose:
        bpy.data.objects.remove(obj, do_unlink=True)
    for action in list(bpy.data.actions):
        bpy.data.actions.remove(action)

    # The rigged character: armature and skinned mesh; helper objects of the GLB go.
    objects = import_glb(source_root / figure['rig'])
    arm = next(o for o in objects if o.type == 'ARMATURE')
    meshes = [o for o in objects if o.type == 'MESH' and o.parent == arm]
    for obj in objects:
        if obj is not arm and obj not in meshes:
            bpy.data.objects.remove(obj, do_unlink=True)
    for action in list(bpy.data.actions):
        bpy.data.actions.remove(action)
    if arm.animation_data:
        arm.animation_data.action = None
    arm.name = figure['id'].capitalize() + 'Rig'
    arm.data.name = arm.name
    body = meshes[0]
    body.name = figure['id'].capitalize() + 'Body'
    body.data.name = body.name

    # Metres: Meshy's 0.01 root scale goes into the vertices and the bone rest pose (clip
    # translations follow in write_clip). The mesh is unparented first, so applying the
    # armature's scale cannot rescale it a second time, and parented again afterwards.
    bpy.ops.object.select_all(action='DESELECT')
    for mesh in meshes:
        mesh.select_set(True)
    bpy.context.view_layer.objects.active = meshes[0]
    bpy.ops.object.parent_clear(type='CLEAR_KEEP_TRANSFORM')
    # Optional uniform scale of the whole figure (figure.json "scale"; U5.5 review: footprint).
    factor = figure.get('scale', 1.0)
    for mesh in meshes:
        mesh.scale = mesh.scale * factor
    arm.scale = arm.scale * factor
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    bpy.ops.object.select_all(action='DESELECT')
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    for mesh in meshes:
        mesh.parent = arm
        mesh.matrix_parent_inverse = Matrix.Identity(4)
    add_root(arm, figure.get('root_bone', 'Root'))

    if not arm.animation_data:
        arm.animation_data_create()
    clips = {}
    for name in CLIP_ORDER:
        spec, samples, frames = sampled[name]
        action = write_clip(arm, name, samples, unit * figure.get('scale', 1.0))
        seconds = (len(samples) - 1) / FPS
        action['loop'] = bool(spec.get('loop', False))
        if 'contact_frame' in spec:
            first, last = frames
            action['contact_seconds'] = round((spec['contact_frame'] - first) / (last - first) * seconds, 4)
        clips[name] = {'source': spec['source'], 'frames': list(frames), 'seconds': round(seconds, 4),
                       'loop': action['loop'], 'contact_seconds': action.get('contact_seconds', -1.0)}
    arm.animation_data.action = None
    for pb in arm.pose.bones:
        pb.matrix_basis = Matrix.Identity(4)

    # One material without image nodes: Unity's Surface material carries the maps (LuneMaterialFactory).
    material = body.material_slots[0].material
    maps = textures(material, str(unity_root / figure['textures']), figure['glow'], figure.get('palette', {}))
    material.name = figure['material']
    tree = material.node_tree
    for node in list(tree.nodes):
        if node.type not in ('BSDF_PRINCIPLED', 'OUTPUT_MATERIAL'):
            tree.nodes.remove(node)
    props = []
    for prop in figure.get('props', []):
        mesh, facts = attach_prop(arm, prop, source_root, unity_root, figure)
        props.append((mesh, prop, facts))
    for image in list(bpy.data.images):
        bpy.data.images.remove(image)
    lod = lane_lod(body, figure['lod']) if 'lod' in figure else None
    for mesh, prop, facts in props:
        if 'lod' in figure:
            facts['lod'] = lane_lod(mesh, dict(figure['lod'], triangles=prop.get('lod_triangles', 600)))
    scene.frame_set(0)
    return {'clips': clips, 'textures': maps, 'unit': unit, 'lod': lod, 'props': [facts for _, _, facts in props]}


def main(build_function, env=None):
    figure_file, source_dir, unity_dir, save_to = options(env)
    result = build_function(figure_file, source_dir, unity_dir)
    print('LUNE_FIGURE', json.dumps(result))
    # Picked up by tools/blender/lune_fbx_export.py for the export report.
    bpy.app.driver_namespace['lune_generator'] = result
    if save_to:
        bpy.ops.wm.save_as_mainfile(filepath=save_to)


# Run as a script or as the exporter's generator; tools/meshy/build_machine.py imports the helpers.
if __name__ in ('__main__', '__lune_generator__'):
    main(build)