"""Export the selected TokenRig GLB as a textured, skinned Unity FBX.

Run in Blender. Does not change the source rig, weights, or prior review assets.
"""
import hashlib
import json
import math
from pathlib import Path

import bpy
import numpy as np
from mathutils import Matrix, Vector
from mathutils.kdtree import KDTree

ROOT = Path(__file__).resolve().parents[2]
SOURCE = ROOT / 'Tools/output/lmgirl_tpose_studio/rig_mixamo41/lmgirl_local.glb'
STATIC = ROOT / 'Tools/output/lmgirl_tpose_studio/static/LmGirlTpose.fbx'
OUT = SOURCE.parent / 'fbx'
FBX = OUT / 'LmGirlTokenRig.fbx'


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def points(mesh):
    return np.array([mesh.matrix_world @ v.co for v in mesh.data.vertices])


def nearest_error(a, b):
    tree = KDTree(len(b))
    for i, p in enumerate(b):
        tree.insert(Vector(p), i)
    tree.balance()
    return max(tree.find(Vector(p))[2] for p in a)


def snapshot(arm, mesh):
    mesh.data.calc_loop_triangles()
    return {
        'points': points(mesh),
        'uv': np.array([list(d.uv) for d in mesh.data.uv_layers.active.data]),
        'triangles': len(mesh.data.loop_triangles),
        'bones': {b.name: {'parent': b.parent.name if b.parent else None,
                         'matrix': np.array(arm.matrix_world @ b.matrix_local)}
                  for b in arm.data.bones},
        'weights': [{mesh.vertex_groups[g.group].name: g.weight for g in v.groups if g.weight > 1e-8}
                    for v in mesh.data.vertices],
    }


def make_material(mesh):
    mat = bpy.data.materials.new('LmGirlTokenRig_Body')
    mat.use_nodes = True
    bsdf = mat.node_tree.nodes.get('Principled BSDF')
    images = {}
    for suffix, socket in [('color', 'Base Color'), ('roughness', 'Roughness'), ('metallic', 'Metallic'), ('normal', None)]:
        path = STATIC.parent / ('LmGirlTpose_' + suffix + '.png')
        image = bpy.data.images.load(str(path), check_existing=True)
        if suffix != 'color':
            image.colorspace_settings.name = 'Non-Color'
        node = mat.node_tree.nodes.new('ShaderNodeTexImage')
        node.image = image
        if suffix == 'normal':
            normal = mat.node_tree.nodes.new('ShaderNodeNormalMap')
            mat.node_tree.links.new(node.outputs['Color'], normal.inputs['Color'])
            mat.node_tree.links.new(normal.outputs['Normal'], bsdf.inputs['Normal'])
        else:
            mat.node_tree.links.new(node.outputs['Color'], bsdf.inputs[socket])
        images[suffix] = str(path.relative_to(ROOT))
    mesh.data.materials.clear()
    mesh.data.materials.append(mat)
    return images


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=str(SOURCE))
    arm = next(o for o in bpy.context.scene.objects if o.type == 'ARMATURE')
    mesh = next(o for o in bpy.context.scene.objects if o.type == 'MESH')
    arm.animation_data_clear()
    assert len(arm.data.bones) == 41
    original = snapshot(arm, mesh)
    print('SOURCE', mesh.name, len(original['points']), original['triangles'],
          original['points'].min(0), original['points'].max(0), 'ARM_MATRIX', arm.matrix_world)

    # The current static FBX has independently reviewed anatomical +Z front.
    # Match its complete imported geometry rather than infer direction from metadata.
    existing = set(bpy.context.scene.objects)
    bpy.ops.import_scene.fbx(filepath=str(STATIC))
    reference = next(o for o in bpy.context.scene.objects if o not in existing and o.type == 'MESH')
    ref = points(reference)
    candidates = []
    for degrees in [0, 180]:
        rotate = Matrix.Rotation(math.radians(degrees), 4, 'Z')
        candidate = np.array([rotate @ Vector(p) for p in original['points']])
        candidates.append((nearest_error(candidate, ref), degrees, rotate))
    error, degrees, conversion = min(candidates, key=lambda item: item[0])
    if error > 1e-5:
        raise RuntimeError('Cannot match reviewed static geometry: ' + str([(e, d) for e, d, _ in candidates]))
    for obj in list(bpy.context.scene.objects):
        if obj not in existing:
            bpy.data.objects.remove(obj, do_unlink=True)
    converted = np.array([conversion @ Vector(p) for p in original['points']])
    lo, hi = converted.min(0), converted.max(0)
    shift = Vector((-(lo[0]+hi[0])/2, -(lo[1]+hi[1])/2, -lo[2]))
    conversion = Matrix.Translation(shift) @ conversion
    matrices = {o: conversion @ o.matrix_world for o in [arm, mesh]}
    mesh.parent = None
    for obj, matrix in matrices.items():
        obj.matrix_world = matrix
    for obj in [arm, mesh]:
        bpy.ops.object.select_all(action='DESELECT')
        obj.select_set(True)
        bpy.context.view_layer.objects.active = obj
        bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    arm.name = 'LmGirlTokenRig'
    mesh.name = 'LmGirlTokenRigMesh'
    texture_paths = make_material(mesh)
    bpy.context.view_layer.update()
    expected = snapshot(arm, mesh)
    assert max(abs(sum(row.values())-1) for row in expected['weights']) < 1e-5
    assert max(len(row) for row in expected['weights']) <= 4
    bpy.ops.object.select_all(action='DESELECT')
    for obj in [arm, mesh]:
        obj.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.export_scene.fbx(
        filepath=str(FBX), use_selection=True, object_types={'ARMATURE', 'MESH'},
        axis_forward='-Z', axis_up='Y', global_scale=1,
        apply_unit_scale=True, apply_scale_options='FBX_SCALE_ALL',
        bake_space_transform=False, add_leaf_bones=False, use_armature_deform_only=False,
        armature_nodetype='NULL', bake_anim=False, mesh_smooth_type='OFF',
        use_mesh_modifiers=False, path_mode='COPY', embed_textures=True)

    # Reimport through Blender's FBX importer, independently of the GLB importer.
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.fbx(filepath=str(FBX))
    actual_arm = next(o for o in bpy.context.scene.objects if o.type == 'ARMATURE')
    actual_mesh = next(o for o in bpy.context.scene.objects if o.type == 'MESH')
    actual = snapshot(actual_arm, actual_mesh)
    assert actual['triangles'] == original['triangles'] == 4999
    assert set(actual['bones']) == set(expected['bones'])
    assert all(actual['bones'][n]['parent'] == expected['bones'][n]['parent'] for n in actual['bones'])
    # FBX may split vertices at corners. Match geometry+UV+weights using nearest positions.
    tree = KDTree(len(expected['points']))
    for i, p in enumerate(expected['points']):
        tree.insert(Vector(p), i)
    tree.balance()
    max_geometry_error = 0.
    max_weight_error = 0.
    mapping = []
    for p, row in zip(actual['points'], actual['weights']):
        _, index, distance = tree.find(Vector(p))
        mapping.append(index)
        max_geometry_error = max(max_geometry_error, distance)
        source_row = expected['weights'][index]
        max_weight_error = max(max_weight_error, max(abs(row.get(n, 0)-source_row.get(n, 0)) for n in set(row)|set(source_row)))
    max_bone_head_error = max(float(np.linalg.norm(actual['bones'][n]['matrix'][:3,3]-expected['bones'][n]['matrix'][:3,3])) for n in actual['bones'])
    # UV distributions are independent of corner ordering. Quantization allows FBX float rounding.
    uv_reference = sorted(map(tuple, np.round(expected['uv'], 5)))
    uv_roundtrip = sorted(map(tuple, np.round(actual['uv'], 5)))
    uv_preserved = uv_reference == uv_roundtrip
    assert max_geometry_error < 1e-5, max_geometry_error
    assert max_weight_error < 1e-5, max_weight_error
    assert max_bone_head_error < 1e-5, max_bone_head_error
    assert uv_preserved
    lo, hi = actual['points'].min(0), actual['points'].max(0)
    report = {
        'source': str(SOURCE.relative_to(ROOT)), 'source_sha256': sha(SOURCE),
        'fbx': str(FBX.relative_to(ROOT)), 'fbx_sha256': sha(FBX),
        'variant': 'TokenRig local constraints, 41 bones, top4; no website weights',
        'sceneUp': '+Y', 'anatomicalFront': '+Z', 'unit': 'meter',
        'conversionAboutBlenderZDegrees': degrees, 'originShiftBlender': list(shift),
        'geometryMatchToReviewedStaticM': error, 'textures': texture_paths,
        'sourceVertices': len(original['points']), 'roundtripVertices': len(actual['points']),
        'triangles': actual['triangles'], 'bones': len(actual['bones']),
        'maxInfluences': max(len(row) for row in actual['weights']),
        'roundtripGeometryErrorM': max_geometry_error, 'roundtripWeightError': max_weight_error,
        'roundtripBoneHeadErrorM': max_bone_head_error, 'uvPreserved': uv_preserved,
        'roundtripBoundsBlender': [lo.tolist(), hi.tolist()],
        'animationIncluded': False,
        'notes': 'Blender workspace is Z-up. Standard skinned FBX axis conversion is retained on the armature node; verify full imported transforms, not raw mesh coordinates alone.'}
    (OUT/'export_verification.json').write_text(json.dumps(report, indent=2), encoding='utf-8')
    print('EXPORTED', json.dumps(report))


if __name__ == '__main__':
    main()