# -*- coding: utf-8 -*-
"""
build_glb_from_fusion_export.py
-------------------------------
Script Blender che ricostruisce un assieme Fusion 360 esportato dallo
script `ExportKinematicGraph.py` (cartella con `hierarchy.json` + `meshes/`)
e produce un file `.glb`.

Uso (da terminale):
    blender --background --python build_glb_from_fusion_export.py -- \
        path/to/export/hierarchy.json path/to/output/plotter.glb

Argomenti dopo il `--`:
    1. Percorso al file `hierarchy.json` esportato da Fusion.
    2. Percorso del file `.glb` di output.

Comportamento:
    * Pulisce la scena Blender.
    * Per ogni nodo della gerarchia crea un Empty (con UUID -> nome univoco).
    * Imposta il parent secondo `parentFullPathName`.
    * Applica la `transform` 4x4 LOCALE letta dal JSON (Fusion = row-major,
      Blender = column-major: facciamo il transpose).
    * Se il nodo ha `meshFile`, importa l'OBJ/STL e lo parenta all'Empty
      con trasformazione locale identita'.
    * Applica lo scaleFactor (cm -> m) al root in modo che l'export GLB
      finisca in metri (convenzione glTF).
    * Esporta in `.glb` (binary glTF 2.0).

Robustezza:
    * Compatibile con Blender 3.x e 4.x (rileva l'importer OBJ corretto).
    * Tutto incapsulato in try/except con log su stderr; exit code != 0
      in caso di errori non recuperabili.
"""

import bpy
import json
import math
import os
import sys
import traceback
from mathutils import Matrix


# =============================================================================
# Argomenti CLI dopo il "--"
# =============================================================================

def _parse_args():
    """Estrae gli argomenti passati dopo il `--` (convenzione Blender)."""
    argv = sys.argv
    if '--' not in argv:
        raise SystemExit(
            'Uso: blender --background --python build_glb_from_fusion_export.py '
            '-- <hierarchy.json> <output.glb>'
        )
    extra = argv[argv.index('--') + 1:]
    if len(extra) < 2:
        raise SystemExit(
            'Servono 2 argomenti dopo "--": hierarchy.json e output.glb'
        )
    return extra[0], extra[1]


# =============================================================================
# Pulizia scena
# =============================================================================

def _wipe_scene():
    """Rimuove tutto dalla scena attiva (oggetti, collezioni orphan, ecc.)."""
    # Selezione di tutto e cancellazione.
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete(use_global=False)

    # Pulizia di datablock orfani (mesh / material / image residui).
    for collection_attr in ('meshes', 'materials', 'images', 'curves',
                            'lights', 'cameras'):
        try:
            data = getattr(bpy.data, collection_attr)
            for item in list(data):
                if item.users == 0:
                    data.remove(item)
        except Exception:
            pass


# =============================================================================
# Conversione matrice JSON -> mathutils.Matrix
# =============================================================================

def _matrix_from_list(values):
    """Converte una lista di 16 float (row-major, come esportato da Fusion)
    in `mathutils.Matrix`. Ritorna l'identita' se `values` non e' valido.
    """
    if not values or len(values) != 16:
        return Matrix.Identity(4)
    rows = [
        values[0:4],
        values[4:8],
        values[8:12],
        values[12:16],
    ]
    return Matrix(rows)


# =============================================================================
# Import mesh (OBJ / STL) cross-version Blender
# =============================================================================

def _import_mesh(filepath):
    """Importa un OBJ o STL e ritorna la lista di Object importati.

    Gestisce sia Blender 4.x (operator nativo `wm.obj_import`) sia 3.x
    (addon `import_scene.obj`). Per STL usa `import_mesh.stl` (presente
    in entrambi).
    """
    before = set(bpy.data.objects)
    ext = os.path.splitext(filepath)[1].lower()

    try:
        if ext == '.obj':
            # Blender 4.x: importer C nativo.
            if hasattr(bpy.ops.wm, 'obj_import'):
                bpy.ops.wm.obj_import(filepath=filepath)
            else:
                # Blender 3.x: addon Python.
                bpy.ops.import_scene.obj(filepath=filepath)
        elif ext == '.stl':
            if hasattr(bpy.ops.wm, 'stl_import'):
                bpy.ops.wm.stl_import(filepath=filepath)
            else:
                bpy.ops.import_mesh.stl(filepath=filepath)
        else:
            print('[WARN] Estensione mesh non supportata: {0}'.format(ext),
                  file=sys.stderr)
            return []
    except Exception:
        print('[ERROR] Import mesh fallito: {0}\n{1}'.format(
            filepath, traceback.format_exc()), file=sys.stderr)
        return []

    return [obj for obj in bpy.data.objects if obj not in before]


# =============================================================================
# Costruzione gerarchia
# =============================================================================

def _make_unique_empty_name(node):
    """Nome univoco per l'Empty di un nodo, usando il fullPathName."""
    # Blender limita i nomi a 63 char; fullPathName puo' essere lungo:
    # usiamo gli ultimi 50 char + un hash breve per evitare collisioni.
    name = node.get('name') or 'node'
    full = node.get('fullPathName') or node.get('id') or name
    suffix = '#{0:08x}'.format(abs(hash(full)) & 0xFFFFFFFF)
    return (name[:50] + suffix)


# Cache materiali condivisi: chiave = tupla RGBA arrotondata -> bpy.Material
_MATERIAL_CACHE = {}


def _make_color_material(color):
    """Ritorna un materiale Blender con `Base Color` = color (RGBA 0-1).

    `color` puo' essere None / lista [r,g,b] / lista [r,g,b,a].
    Il materiale viene cached per RGBA arrotondata, in modo che colori
    uguali condividano lo stesso material slot.
    Ritorna None se `color` non e' valido.
    """
    if not color or not isinstance(color, list) or len(color) < 3:
        return None
    try:
        r, g, b = float(color[0]), float(color[1]), float(color[2])
        a = float(color[3]) if len(color) >= 4 else 1.0
    except Exception:
        return None

    # Chiave cache: arrotonda a 3 decimali (~256 livelli per canale).
    key = (round(r, 3), round(g, 3), round(b, 3), round(a, 3))
    cached = _MATERIAL_CACHE.get(key)
    if cached is not None:
        return cached

    mat = bpy.data.materials.new(name='FusionMat_{0:02x}{1:02x}{2:02x}'.format(
        int(r * 255), int(g * 255), int(b * 255)))
    mat.use_nodes = True
    try:
        bsdf = mat.node_tree.nodes.get('Principled BSDF')
        if bsdf is not None:
            bsdf.inputs['Base Color'].default_value = (r, g, b, a)
            # Se c'e' trasparenza, attiviamo l'alpha blending.
            if a < 1.0:
                bsdf.inputs['Alpha'].default_value = a
                mat.blend_method = 'BLEND'
    except Exception:
        # Fallback: colore di viewport (sempre disponibile).
        try:
            mat.diffuse_color = (r, g, b, a)
        except Exception:
            pass

    # Anche il colore "viewport" (per Solid view e fallback).
    try:
        mat.diffuse_color = (r, g, b, a)
    except Exception:
        pass

    _MATERIAL_CACHE[key] = mat
    return mat


def build_scene(hierarchy_path):
    """Costruisce la scena Blender a partire da hierarchy.json.

    Ritorna l'Empty 'root' a cui tutto il modello e' parentato.
    """
    with open(hierarchy_path, 'r', encoding='utf-8') as f:
        payload = json.load(f)

    nodes = payload.get('nodes', [])
    metadata = payload.get('metadata', {}) or {}
    scale_factor = float(metadata.get('scaleFactor') or 0.01)
    base_dir = os.path.dirname(os.path.abspath(hierarchy_path))

    # Indice: id (== fullPathName, vuoto per il root) -> Empty Blender.
    id_to_empty = {}

    # 1) Pass 1: crea tutti gli Empty senza parenting.
    for node in nodes:
        node_id = node.get('id')
        if node_id is None:
            continue

        empty = bpy.data.objects.new(_make_unique_empty_name(node), None)
        empty.empty_display_type = 'ARROWS'
        empty.empty_display_size = 0.1
        bpy.context.collection.objects.link(empty)

        # Salviamo metadati utili come custom properties.
        # I campi `fusion_name` e `fusion_path` sono il contratto ufficiale col
        # viewer Three.js (vedi FUSION_GLB_CONTRACT.md §3.2): devono coincidere
        # byte-per-byte con `Joint.child` / `Joint.childFullPath` del joints.json.
        try:
            occ_name = node.get('name') or ''
            full_path = node.get('fullPathName') or ''
            empty['fusion_name']         = occ_name
            empty['fusion_path']         = full_path
            # Alias storici (back-compat, possono essere rimossi in futuro).
            empty['fusion_id']           = node_id
            empty['fusion_fullPathName'] = full_path
            empty['fusion_componentName'] = node.get('componentName') or ''
        except Exception:
            pass

        id_to_empty[node_id] = empty

    # 2) Pass 2: SOLO parenting + import mesh. Il transform viene applicato
    # dopo in un BFS top-down (vedi pass 3).
    root_node_id = '__ROOT__'
    for node in nodes:
        node_id = node.get('id')
        empty = id_to_empty.get(node_id)
        if empty is None:
            continue

        # Parent:
        #   parentFullPathName == None -> nessun parent (e' il root)
        #   parentFullPathName == ''   -> figlio del root (export vecchio)
        #   altrimenti                 -> id del parent
        parent_path = node.get('parentFullPathName')
        if parent_path is not None and node_id != root_node_id:
            # Retrocompatibilita': stringa vuota = root.
            if parent_path == '':
                parent_path = root_node_id
            parent_empty = id_to_empty.get(parent_path)
            if parent_empty is None:
                # Orfano: parent non trovato, agganciamo al root se esiste.
                parent_empty = id_to_empty.get(root_node_id)
                if parent_empty is not None:
                    print('[WARN] Parent non trovato per "{0}" (cercavo "{1}'
                          '"). Aggancio al root.'.format(node_id, parent_path),
                          file=sys.stderr)
            if parent_empty is not None and parent_empty is not empty:
                empty.parent = parent_empty

        # Mesh associata. Parentata all'Empty con identita': la posizione
        # finale verra' dall'Empty (gestita nel pass 3).
        mesh_rel = node.get('meshFile')
        if mesh_rel:
            mesh_abs = os.path.join(base_dir, mesh_rel)
            if not os.path.isfile(mesh_abs):
                print('[WARN] Mesh non trovata: {0}'.format(mesh_abs),
                      file=sys.stderr)
                continue
            imported = _import_mesh(mesh_abs)
            # Materiale colorato dal `color` del nodo (se presente). I
            # materiali sono cached per RGBA cosi' colori uguali condividono
            # lo stesso material slot nel GLB (output piu' piccolo).
            mat_color = _make_color_material(node.get('color'))
            for obj in imported:
                obj.parent = empty
                obj.matrix_local = Matrix.Identity(4)
                if mat_color is not None:
                    try:
                        # Sovrascriviamo tutti gli slot esistenti per essere
                        # sicuri che il colore Fusion venga rispettato.
                        if obj.type == 'MESH':
                            obj.data.materials.clear()
                            obj.data.materials.append(mat_color)
                    except Exception:
                        pass

    # 3) Trovare il nodo "root" e applicare lo scaleFactor cm -> m.
    root_empty = id_to_empty.get(root_node_id)
    if root_empty is None:
        # Fallback: prendiamo qualsiasi Empty senza parent.
        for e in id_to_empty.values():
            if e.parent is None:
                root_empty = e
                break

    if root_empty is not None and scale_factor and scale_factor != 1.0:
        # Scala uniforme applicata SOLO al root: i figli ereditano la scala.
        root_empty.scale = (scale_factor, scale_factor, scale_factor)

    # Forziamo l'update della scene cosi' matrix_world del root e' calcolato
    # prima di propagare le world matrices dei figli.
    bpy.context.view_layer.update()

    # 4) Pass 4: BFS top-down e assegnazione di matrix_world.
    # IMPORTANTE: `Occurrence.transform2` di Fusion ritorna la matrice in
    # coordinate WORLD/ROOT (NON relative al parent immediato). Quindi NON
    # possiamo settare `matrix_local` (darebbe traslazioni duplicate per i
    # nodi dentro sub-assembly non a origine). Settiamo `matrix_world` e
    # Blender calcola da solo il local relativo al parent.
    # La matrice esportata e' in cm Fusion: la convertiamo in metri
    # moltiplicando per scale_root (che include lo scaleFactor).
    from collections import deque
    scale_root_mat = Matrix.Scale(scale_factor, 4) if scale_factor else Matrix.Identity(4)

    id_to_node = {n.get('id'): n for n in nodes}
    children_map = {}
    for n in nodes:
        p = n.get('parentFullPathName')
        if p is None:
            continue
        if p == '':
            p = root_node_id
        children_map.setdefault(p, []).append(n.get('id'))

    visited = {root_node_id}
    queue = deque([root_node_id])
    while queue:
        cur_id = queue.popleft()
        for child_id in children_map.get(cur_id, []):
            if child_id in visited:
                continue
            visited.add(child_id)
            queue.append(child_id)
            child_empty = id_to_empty.get(child_id)
            child_node = id_to_node.get(child_id)
            if child_empty is None or child_node is None:
                continue
            mat_world_cm = _matrix_from_list(child_node.get('transform'))
            # world finale (m) = scale_root @ world_fusion (cm)
            child_empty.matrix_world = scale_root_mat @ mat_world_cm

    return root_empty


# =============================================================================
# Joints / Motion Links -> glTF scene.extras
# =============================================================================

def embed_joints(joints_json_path):
    """Carica joints.json e lo salva come custom property della Scene.

    Con `export_extras=True` Blender mette questo dict in `scene.extras`
    del glTF, quindi in Three.js puoi leggerlo via
    `gltf.parser.json.scenes[i].extras.fusion`.

    Inoltre converte le coordinate da Fusion (cm) a metri (m) cosi' i
    valori sono coerenti col modello esportato col `scaleFactor`.
    """
    if not joints_json_path or not os.path.isfile(joints_json_path):
        print('[INFO] joints.json non trovato, skip embedding joints/motion.',
              file=sys.stderr)
        return

    try:
        with open(joints_json_path, 'r', encoding='utf-8') as f:
            payload = json.load(f)
    except Exception:
        print('[WARN] joints.json non leggibile:\n' + traceback.format_exc(),
              file=sys.stderr)
        return

    # Scaling cm -> m sulle posizioni (origin) e sui valori dei limiti
    # lineari (slideLimits / slideValue). Gli assi sono direzioni
    # unitarie -> non vanno scalati. Le rotazioni sono in radianti
    # -> non vanno scalate.
    metadata = payload.get('metadata', {}) or {}
    scale = float(metadata.get('scaleFactor') or 0.01)

    def _scale_point(p):
        if p is None or not isinstance(p, list):
            return p
        return [v * scale for v in p]

    def _scale_limits(lim, is_linear):
        if lim is None or not is_linear:
            return lim
        out = dict(lim)
        for k in ('minimumValue', 'maximumValue', 'restValue'):
            if out.get(k) is not None:
                try:
                    out[k] = out[k] * scale
                except Exception:
                    pass
        return out

    def _scale_joint(j):
        j = dict(j)
        for k in ('origin', 'originTwo'):
            if k in j:
                j[k] = _scale_point(j.get(k))
        if 'slideLimits' in j:
            j['slideLimits'] = _scale_limits(j.get('slideLimits'), True)
        if 'slideValue' in j and j.get('slideValue') is not None:
            try:
                j['slideValue'] = j['slideValue'] * scale
            except Exception:
                pass
        # rotationLimits / rotationValue restano invariati (radianti).
        return j

    fusion_data = {
        'metadata':      metadata,
        'joints':        [_scale_joint(j) for j in payload.get('joints', [])],
        'asBuiltJoints': [_scale_joint(j) for j in payload.get('asBuiltJoints', [])],
        'motionLinks':   payload.get('motionLinks', []),
        'rigidGroups':   payload.get('rigidGroups', []),
    }

    # Custom property sulla Scene -> finisce in scene.extras del glTF.
    scene = bpy.context.scene
    try:
        # Serializziamo in stringa JSON: glTF custom properties supportano
        # primitivi/array/dict ma alcuni viewer non gestiscono dict annidati
        # profondi. Una stringa JSON e' sempre safe da parsare in Three.js.
        scene['fusion'] = json.dumps(fusion_data, ensure_ascii=False)
        print('[OK] Joints embedded: joints={0} as_built={1} '
              'motion_links={2} rigid_groups={3}'.format(
                  len(fusion_data['joints']),
                  len(fusion_data['asBuiltJoints']),
                  len(fusion_data['motionLinks']),
                  len(fusion_data['rigidGroups'])))
    except Exception:
        print('[WARN] Embed joints fallito:\n' + traceback.format_exc(),
              file=sys.stderr)


# =============================================================================
# Export GLB
# =============================================================================

def export_glb(output_path):
    """Esporta la scena corrente in formato GLB (glTF 2.0 binario)."""
    out_dir = os.path.dirname(os.path.abspath(output_path))
    if out_dir and not os.path.isdir(out_dir):
        os.makedirs(out_dir, exist_ok=True)

    bpy.ops.export_scene.gltf(
        filepath=output_path,
        export_format='GLB',
        export_apply=True,          # applica modifiers
        export_yup=True,            # convenzione glTF: Y-up
        use_visible=False,
        use_selection=False,
        export_extras=True,         # custom properties -> extras glTF
        export_materials='EXPORT',  # esporta i materiali Principled BSDF
    )


# =============================================================================
# Main
# =============================================================================

def main():
    try:
        hierarchy_path, output_path = _parse_args()

        if not os.path.isfile(hierarchy_path):
            raise SystemExit(
                'File hierarchy non trovato: {0}'.format(hierarchy_path))

        print('[INFO] Hierarchy: {0}'.format(hierarchy_path))
        print('[INFO] Output:    {0}'.format(output_path))

        _wipe_scene()
        root = build_scene(hierarchy_path)
        if root is None:
            print('[WARN] Nessun nodo root identificato.', file=sys.stderr)

        # Cerca joints.json accanto al hierarchy e lo incorpora nel GLB
        # come scene.extras.fusion (stringa JSON).
        joints_path = os.path.join(
            os.path.dirname(os.path.abspath(hierarchy_path)), 'joints.json')
        embed_joints(joints_path)

        export_glb(output_path)

        print('[OK] Esportato: {0}'.format(output_path))

    except SystemExit:
        raise
    except Exception:
        print('[ERROR] Build GLB fallito:\n' + traceback.format_exc(),
              file=sys.stderr)
        raise SystemExit(1)


if __name__ == '__main__':
    main()