"""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

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


def options():
    """Figure file and roots: from the exporter (globals) or the command line."""
    g = 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(__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 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)
        sampled[name] = (spec, sample_clip(arm_anim, source, (first, last), seconds), (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')
    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)
        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)
    for image in list(bpy.data.images):
        bpy.data.images.remove(image)
    scene.frame_set(0)
    return {'clips': clips, 'textures': maps, 'unit': unit}


figure_file, source_dir, unity_dir, save_to = options()
result = build(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)