"""Runs INSIDE Blender. Rebuilds a Unity prefab from its source FBX models and exports one GLB.

    blender --background --python tools/blender_build.py -- spec.json

The spec (written by convert_assets.py) lists every renderer the prefab places: which FBX and
which mesh inside it, the resolved Unity material, and the world transform. This script imports
each FBX once, then instances the mesh datablock per renderer so a wall repeated 180 times costs
one mesh, not 180.

Coordinates: Blender's FBX importer lands Unity's models at (x, z, y) — Y and Z swapped — and the
glTF exporter maps Blender (x, y, z) back to (x, z, -y). The round trip therefore delivers
(unity_x, unity_y, -unity_z), which is exactly the convention Arena.ts already uses for depth. So
transforms are converted into Blender's basis here and need no fixing up on the JavaScript side.
"""
import json
import os
import sys

import math
import bpy
from mathutils import Matrix, Quaternion, Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from fbx_ascii import read_ascii_fbx  # noqa: E402

# ---------------------------------------------------------------- coordinates

def to_blender_pos(p):
    return Vector((p[0], p[2], p[1]))


def to_blender_quat(q):
    """Swapping two axes is a reflection, so conjugating by it reverses the sense of rotation:
    the axis maps (x, y, z) -> (x, z, y) and the angle negates."""
    return Quaternion((q[3], -q[0], -q[2], -q[1]))


def to_blender_scale(s):
    return Vector((s[0], s[2], s[1]))


# The same axis swap, as a matrix, for baking into mesh data. Its determinant is -1 — Unity is
# left-handed and Blender is not — so every mesh it touches needs its winding reversed afterwards
# or the faces end up inside out.
UNITY_TO_BLENDER = Matrix(((1, 0, 0, 0),
                           (0, 0, 1, 0),
                           (0, 1, 0, 0),
                           (0, 0, 0, 1)))


def permutation_matrix(perm, signs):
    """Raw FBX axes -> Unity mesh axes. `unity[i] = signs[i] * raw[perm[i]]`.

    The signs are not decoration: Unity negates X converting a right-handed FBX into its own
    left-handed space, so a transform without them mirrors every mesh. That is invisible on a
    symmetric wall panel and very visible once those panels are assembled into a structure.
    """
    m = Matrix.Identity(4)
    for i in range(3):
        for j in range(3):
            m[i][j] = float(signs[i]) if j == perm[i] else 0.0
    return m


def geometric_matrix(g):
    """An FBX Model node's geometric transform as a 4x4, in the file's own units.

    Scaling, then an XYZ Euler rotation (Rz*Ry*Rx, degrees), then translation — the FBX SDK's
    definition, and the one Unity bakes into the mesh. This is what the raw vertex array is
    missing: Blender folds it into the object matrix, which this build deliberately discards.
    """
    if not g:
        return Matrix.Identity(4)
    sc = g.get('GeometricScaling', (1.0, 1.0, 1.0))
    rot = g.get('GeometricRotation', (0.0, 0.0, 0.0))
    tr = g.get('GeometricTranslation', (0.0, 0.0, 0.0))
    m = Matrix.Diagonal((sc[0], sc[1], sc[2], 1.0))
    rx = Matrix.Rotation(math.radians(rot[0]), 4, 'X')
    ry = Matrix.Rotation(math.radians(rot[1]), 4, 'Y')
    rz = Matrix.Rotation(math.radians(rot[2]), 4, 'Z')
    return Matrix.Translation(tr) @ rz @ ry @ rx @ m


def bake_to_blender(mesh, matrix_world, perm, signs, geometric=None):
    """Put a mesh into Blender's basis, from the raw FBX vertices, once and for all.

    Three things happen here, and all three were learned the hard way:

    * Only the import SCALE is taken from the object. An FBX node's own transform becomes a
      GameObject Transform on Unity's side, and the prefab walk already carries that Transform in
      the matrix each instance is placed with; baking it here too applies it twice.
    * `perm` corrects the raw FBX axes to the ones Unity actually stores. It is not the same for
      every asset pack and cannot be read off the file's declared up-axis — `fbx_axis.py` derives it
      from Unity's own BoxCollider sizes. Without it one 3ds Max kit's models, which were the whole
      arena, arrive rotated: floors stand on edge and walls lie flat.
    * `UNITY_TO_BLENDER` then swaps into Blender's basis, which mirrors, so the winding is reversed.
    """
    sc = matrix_world.to_scale()
    # Blender folds the node's GeometricScaling into the object scale it imports, so `sc` is
    # (unit scale x GeometricScaling). The geometric matrix below applies that scaling itself —
    # taking `sc` as-is applied it twice, which is why the centre statue (GeometricScaling 0.28,
    # a statue pack) came out at a quarter of its size while the kit pieces, whose
    # GeometricScaling is 1, were unaffected. Divide it back out: what remains is the unit scale
    # Unity applies with useFileScale (2.54 -> 0.0254 for that file), and the node's own Lcl
    # scaling, which is 1 across these packs and which Unity ignores for a mesh anyway.
    if geometric:
        gs = geometric.get('GeometricScaling', (1.0, 1.0, 1.0))
        sc = Vector((sc.x / gs[0] if abs(gs[0]) > 1e-9 else sc.x,
                     sc.y / gs[1] if abs(gs[1]) > 1e-9 else sc.y,
                     sc.z / gs[2] if abs(gs[2]) > 1e-9 else sc.z))
    # Geometric transform is in file units, so it goes on BEFORE the import scale.
    total = (UNITY_TO_BLENDER @ permutation_matrix(perm, signs)
             @ Matrix.Diagonal((sc.x, sc.y, sc.z, 1.0)) @ geometric_matrix(geometric))
    mesh.transform(total)
    # Winding reverses once per reflection. UNITY_TO_BLENDER is one; a negated axis in the
    # permutation is another, and two of them cancel — so flip only when the total is a mirror.
    if total.to_3x3().determinant() < 0:
        mesh.flip_normals()


# ---------------------------------------------------------------- textures

# The source art is 2048x2048 TIFF, which is far more than a web build can carry and far more
# than a top-down game shows: the level is lit by one moving point light, most surfaces are seen at a
# distance, and the whole thing renders in the dark. Masks carry no visible detail of their own,
# so they drop further than colour and normals do.
MAX_TEX = {'baseColor': 512, 'normal': 512, 'mask': 256, 'emissive': 256}

_img_cache = {}


def load_image(path, non_color, slot=None):
    key = (path, non_color)
    if key in _img_cache:
        return _img_cache[key]
    try:
        img = bpy.data.images.load(path, check_existing=True)
    except RuntimeError:
        _img_cache[key] = None
        return None
    if non_color:
        img.colorspace_settings.name = 'Non-Color'
    limit = MAX_TEX.get(slot or '', 512)
    w, h = img.size
    if w > limit or h > limit:
        img.scale(min(w, limit), min(h, limit))
    _img_cache[key] = img
    return img


_orm_cache = {}


def repack_mask_to_orm(path):
    """Unity packs metallic in R and smoothness in A; glTF wants roughness in G, metallic in B.

    Rewriting the channels here means the exporter can point Roughness and Metallic straight at
    one image, which is what glTF stores natively — no bake step, no guessing on import.
    """
    if path in _orm_cache:
        return _orm_cache[path]
    src = load_image(path, True, 'mask')
    if src is None:
        _orm_cache[path] = None
        return None
    import numpy as np
    w, h = src.size
    if w == 0 or h == 0:
        _orm_cache[path] = None
        return None
    buf = np.empty(w * h * 4, dtype=np.float32)
    src.pixels.foreach_get(buf)
    buf = buf.reshape(-1, 4)
    out = np.ones_like(buf)
    out[:, 0] = buf[:, 1]          # AO, where Unity's mask keeps it
    out[:, 1] = 1.0 - buf[:, 3]    # roughness = 1 - smoothness
    out[:, 2] = buf[:, 0]          # metallic
    orm = bpy.data.images.new(os.path.basename(path) + '_orm', w, h, alpha=True, float_buffer=False)
    orm.colorspace_settings.name = 'Non-Color'
    orm.pixels.foreach_set(out.reshape(-1))
    orm.pack()
    _orm_cache[path] = orm
    return orm


# ---------------------------------------------------------------- materials

_mat_cache = {}


def build_material(spec):
    name = spec['name']
    if name in _mat_cache:
        return _mat_cache[name]

    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    nt = mat.node_tree
    bsdf = nt.nodes['Principled BSDF']
    tex = spec.get('textures', {})
    cols = spec.get('colors', {})
    floats = spec.get('floats', {})

    def place(node, x, y):
        node.location = (x, y)
        return node

    base = cols.get('baseColor')
    if base:
        bsdf.inputs['Base Color'].default_value = (base[0], base[1], base[2], 1.0)

    if tex.get('baseColor'):
        img = load_image(tex['baseColor'], False, 'baseColor')
        if img:
            n = place(nt.nodes.new('ShaderNodeTexImage'), -700, 300)
            n.image = img
            if base and (base[0] < 0.99 or base[1] < 0.99 or base[2] < 0.99):
                # Unity multiplies the map by the base colour; keep that tint.
                mix = place(nt.nodes.new('ShaderNodeMixRGB'), -400, 300)
                mix.blend_type = 'MULTIPLY'
                mix.inputs['Fac'].default_value = 1.0
                mix.inputs['Color2'].default_value = (base[0], base[1], base[2], 1.0)
                nt.links.new(n.outputs['Color'], mix.inputs['Color1'])
                nt.links.new(mix.outputs['Color'], bsdf.inputs['Base Color'])
            else:
                nt.links.new(n.outputs['Color'], bsdf.inputs['Base Color'])

    if tex.get('normal'):
        img = load_image(tex['normal'], True, 'normal')
        if img:
            n = place(nt.nodes.new('ShaderNodeTexImage'), -700, -100)
            n.image = img
            nm = place(nt.nodes.new('ShaderNodeNormalMap'), -400, -100)
            nt.links.new(n.outputs['Color'], nm.inputs['Color'])
            nt.links.new(nm.outputs['Normal'], bsdf.inputs['Normal'])

    if tex.get('mask'):
        orm = repack_mask_to_orm(tex['mask'])
        if orm:
            n = place(nt.nodes.new('ShaderNodeTexImage'), -700, -450)
            n.image = orm
            sep = place(nt.nodes.new('ShaderNodeSeparateColor'), -400, -450)
            nt.links.new(n.outputs['Color'], sep.inputs['Color'])
            nt.links.new(sep.outputs['Green'], bsdf.inputs['Roughness'])
            nt.links.new(sep.outputs['Blue'], bsdf.inputs['Metallic'])
    else:
        if 'metallic' in floats:
            bsdf.inputs['Metallic'].default_value = floats['metallic']
        if 'smoothness' in floats:
            bsdf.inputs['Roughness'].default_value = 1.0 - floats['smoothness']

    # Only light a material up when Unity actually had emission switched on. These shaders keep a
    # colour in the emissive slot regardless, so trusting the value alone turns most of the level
    # into a white light source.
    emis = cols.get('emissive') if spec.get('emissiveEnabled') else None
    if tex.get('emissive'):
        img = load_image(tex['emissive'], False, 'emissive')
        if img:
            n = place(nt.nodes.new('ShaderNodeTexImage'), -700, 650)
            n.image = img
            nt.links.new(n.outputs['Color'], bsdf.inputs['Emission Color'])
            bsdf.inputs['Emission Strength'].default_value = 1.0
    elif emis and max(emis[:3]) > 0.01:
        bsdf.inputs['Emission Color'].default_value = (emis[0], emis[1], emis[2], 1.0)
        bsdf.inputs['Emission Strength'].default_value = 1.0

    _mat_cache[name] = mat
    return mat


# ---------------------------------------------------------------- meshes

_fbx_cache = {}


def is_ascii_fbx(path):
    with open(path, 'rb') as fh:
        return b'Kaydara' not in fh.read(32)


def build_ascii_meshes(path):
    """Build mesh datablocks from an ASCII FBX, which Blender's importer will not read.

    Axes and scale are matched to what the binary importer does — Blender (x, y, z) takes FBX
    (x, z, y), scaled by UnitScaleFactor/100 — so ASCII and binary models land in the same space.
    """
    try:
        meshes, unit = read_ascii_fbx(path)
    except Exception as exc:  # noqa: BLE001
        print(f'FBXFAIL {path}: {exc}')
        return {}

    scale = unit / 100.0
    out = {}
    for src in meshes:
        # Same as the binary path: Unity negates X on import, then Blender's basis swaps Y/Z.
        verts = [(-v[0] * scale, v[2] * scale, v[1] * scale) for v in src['verts']]
        polys = [p for p in src['polys'] if len(p) >= 3 and max(p) < len(verts)]
        if not verts or not polys:
            continue

        mesh = bpy.data.meshes.new(src['name'])
        mesh.from_pydata(verts, [], polys)
        mesh.validate()
        # verts were swapped from Unity axes above, which mirrors them.
        mesh.flip_normals()

        uvs = src['uvs']
        if uvs:
            layer = mesh.uv_layers.new(name='UVMap')
            index = src['uvIndex']
            direct = src['uvReference'] != 'IndexToDirect'
            loop = 0
            for poly in mesh.polygons:
                for _ in poly.loop_indices:
                    at = loop if direct else (index[loop] if loop < len(index) else 0)
                    if 0 <= at * 2 + 1 < len(uvs):
                        layer.data[loop].uv = (uvs[at * 2], uvs[at * 2 + 1])
                    loop += 1

        normals = src['normals']
        if normals and src['normalMapping'] == 'ByPolygonVertex' and len(normals) >= len(mesh.loops) * 3:
            try:
                mesh.normals_split_custom_set([
                    (normals[i * 3], normals[i * 3 + 2], normals[i * 3 + 1])
                    for i in range(len(mesh.loops))
                ])
                for poly in mesh.polygons:
                    poly.use_smooth = True
            except (RuntimeError, TypeError) as exc:
                print(f'NORMALS {path}: {exc}')

        out[src['name']] = mesh
    return out


def import_fbx(path, perm=(0, 1, 2), signs=(-1, 1, 1), geometric=None):
    """Import once and hand back {mesh name: mesh datablock}, in final local units."""
    if path in _fbx_cache:
        return _fbx_cache[path]
    if is_ascii_fbx(path):
        built = build_ascii_meshes(path)
        _fbx_cache[path] = built
        return built
    before = set(bpy.data.objects)
    try:
        bpy.ops.import_scene.fbx(filepath=path)
    except Exception as exc:  # noqa: BLE001 - a bad model should not sink the whole prefab
        print(f'FBXFAIL {path}: {exc}')
        _fbx_cache[path] = {}
        return {}
    new = [o for o in bpy.data.objects if o not in before]

    meshes = {}
    baked = set()
    mesh_index = 0   # pairs with read_fbx's per-mesh geometric list, both in file order
    for o in new:
        if o.type != 'MESH':
            continue
        if o.data.name not in baked:
            gi = geometric[mesh_index] if (geometric and mesh_index < len(geometric)) else None
            bake_to_blender(o.data, o.matrix_world, perm, signs, gi)
            mesh_index += 1
            baked.add(o.data.name)
        meshes[o.name] = o.data
    for o in new:
        bpy.data.objects.remove(o, do_unlink=True)
    _fbx_cache[path] = meshes
    return meshes


_variant_cache = {}


def variant_for(data, mats):
    """Material slots live on the mesh, and instances share it — so one mesh worn two different
    ways needs two datablocks. Everything using the same combination still shares one."""
    key = (data.name, tuple(m['name'] for m in mats))
    hit = _variant_cache.get(key)
    if hit is not None:
        return hit
    if not mats:
        _variant_cache[key] = data
        return data
    if _variant_cache and any(k[0] == data.name for k in _variant_cache):
        data = data.copy()
    data.materials.clear()
    for m in mats:
        data.materials.append(build_material(m))
    _variant_cache[key] = data
    return data


def pick_mesh(meshes, wanted):
    if not meshes:
        return None
    if wanted and wanted in meshes:
        return meshes[wanted]
    if wanted:
        for k, v in meshes.items():
            if k.split('.')[0] == wanted:
                return v
    if len(meshes) == 1:
        return next(iter(meshes.values()))
    return next(iter(meshes.values()))


# ---------------------------------------------------------------- build

def main():
    argv = sys.argv[sys.argv.index('--') + 1:]
    spec_path = argv[0]
    with open(spec_path) as fh:
        spec = json.load(fh)

    bpy.ops.wm.read_factory_settings(use_empty=True)
    coll = bpy.context.scene.collection

    materials = {m['path']: m for m in spec.get('materials', [])}
    placed = 0
    skipped = 0

    axis = spec.get('axis', {})
    for r in spec['renderers']:
        ax = axis.get(r['fbx']) or {}
        meshes = import_fbx(r['fbx'],
                            tuple(ax.get('perm', (0, 1, 2))),
                            tuple(ax.get('signs', (-1, 1, 1))),
                            spec.get('geometric', {}).get(r['fbx']))
        data = pick_mesh(meshes, r.get('meshName'))
        if data is None:
            skipped += 1
            continue
        mats = [materials.get(p) for p in r.get('materials', [])]
        mats = [m for m in mats if m]
        data = variant_for(data, mats)

        obj = bpy.data.objects.new(r['name'][:60], data)
        coll.objects.link(obj)
        obj.rotation_mode = 'QUATERNION'
        obj.location = to_blender_pos(r['pos'])
        obj.rotation_quaternion = to_blender_quat(r['quat'])
        obj.scale = to_blender_scale(r['scale'])
        placed += 1

    print(f'PLACED {placed} SKIPPED {skipped}')

    out = spec['out']
    os.makedirs(os.path.dirname(out), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=out,
        export_format='GLB',
        export_apply=False,
        export_materials='EXPORT',
        export_image_format='JPEG',
        export_jpeg_quality=80,
        export_yup=True,
    )
    print(f'EXPORTED {out} {os.path.getsize(out)}')


main()