# SketchUp Tag glTF Importer
# Imports all .glb files from a folder (one file per SketchUp Tag),
# producing one mesh object per file inside a timestamped collection.
#
# v4.2.0 - Materials and textures reuse existing datablocks by name across
#   import runs, so multiple models in one .blend share one "Steel" etc.
#
# v4.1.0 - Direct GLB reader:
#   SketchUp glTF exports carry one node per group/component (a real project
#   tag can hold thousands). Blender's stock glTF importer creates an object
#   per node, and merging those afterward is unusably slow at project scale.
#   This version parses the .glb directly with numpy and builds exactly ONE
#   mesh object per file. No intermediate objects, no joins, no empties.
#   Materials (base color + embedded textures) and UVs are preserved.
#   Custom normals are not; flat-shaded architecture doesn't need them.
#   If direct parsing fails on a file (e.g. Draco-compressed), it falls back
#   to Blender's stock glTF importer for that file only.
#
# v4.0.x history: stock importer + join approach. Superseded.

bl_info = {
    "name": "SketchUp to Blender",
    "author": "Spencer Clem",
    "version": (4, 2, 0),
    "blender": (4, 2, 0),
    "location": "3D Viewport > Sidebar > SKP Import",
    "description": "Fast batch import of per-tag .glb files, one object per tag",
    "category": "Import-Export",
}

import os
import json
import struct
import datetime
import tempfile

import numpy as np
import bpy
from bpy.props import StringProperty, FloatProperty
from bpy.types import Operator, Panel, PropertyGroup

GLTF_EXTENSIONS = (".glb",)

# ----------------------------
# GLB parsing (no bpy in this section)
# ----------------------------

COMP_DTYPE = {5120: np.int8, 5121: np.uint8, 5122: np.int16,
              5123: np.uint16, 5125: np.uint32, 5126: np.float32}
TYPE_SIZE = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT4": 16}


def parse_glb(path):
    """Split a .glb into its JSON header and binary buffer."""
    with open(path, "rb") as f:
        data = f.read()
    magic, _version, length = struct.unpack_from("<III", data, 0)
    if magic != 0x46546C67:
        raise ValueError("Not a GLB file")
    offset, gltf, binbuf = 12, None, None
    while offset < length:
        clen, ctype = struct.unpack_from("<II", data, offset)
        offset += 8
        chunk = data[offset:offset + clen]
        if ctype == 0x4E4F534A:
            gltf = json.loads(chunk)
        elif ctype == 0x004E4942:
            binbuf = chunk
        offset += clen
    if gltf is None or binbuf is None:
        raise ValueError("GLB missing JSON or BIN chunk")
    if gltf.get("extensionsRequired"):
        raise ValueError(f"Unsupported required extensions: {gltf['extensionsRequired']}")
    return gltf, binbuf


def read_accessor(gltf, binbuf, idx):
    acc = gltf["accessors"][idx]
    if "sparse" in acc:
        raise ValueError("Sparse accessors not supported")
    bv = gltf["bufferViews"][acc["bufferView"]]
    dtype = COMP_DTYPE[acc["componentType"]]
    ncomp = TYPE_SIZE[acc["type"]]
    itemsize = np.dtype(dtype).itemsize
    start = bv.get("byteOffset", 0) + acc.get("byteOffset", 0)
    stride = bv.get("byteStride")
    count = acc["count"]
    if stride and stride != ncomp * itemsize:
        raw = np.frombuffer(binbuf, dtype=np.uint8, count=stride * count, offset=start)
        raw = raw.reshape(count, stride)[:, :ncomp * itemsize]
        arr = raw.copy().view(dtype).reshape(count, ncomp)
    else:
        arr = np.frombuffer(binbuf, dtype=dtype, count=count * ncomp, offset=start)
        arr = arr.reshape(count, ncomp) if ncomp > 1 else arr
    return arr


def trs_matrix(node):
    m = np.eye(4)
    if "matrix" in node:
        return np.array(node["matrix"], dtype=np.float64).reshape(4, 4).T
    if "scale" in node:
        m[:3, :3] *= np.array(node["scale"])
    if "rotation" in node:
        x, y, z, w = node["rotation"]
        r = np.array([
            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]])
        m[:3, :3] = r @ m[:3, :3]
    if "translation" in node:
        m[:3, 3] = node["translation"]
    return m


def collect_geometry(gltf, binbuf):
    """Walk the scene graph, return combined arrays in Blender coordinates.

    Returns (verts Nx3, tris Mx3, tri_material M, loop_uvs Mx3x2 flattened)
    plus the list of glTF material indices actually used, in slot order.
    """
    nodes = gltf["nodes"]
    verts_list, tris_list, mat_list, uv_list = [], [], [], []
    v_offset = 0
    scene = gltf["scenes"][gltf.get("scene", 0)]
    stack = [(i, np.eye(4)) for i in scene["nodes"]]
    while stack:
        ni, parent = stack.pop()
        node = nodes[ni]
        world = parent @ trs_matrix(node)
        for child in node.get("children", []):
            stack.append((child, world))
        if "mesh" not in node:
            continue
        for prim in gltf["meshes"][node["mesh"]]["primitives"]:
            if prim.get("mode", 4) != 4:
                continue  # points/lines: not relevant to this pipeline
            pos = read_accessor(gltf, binbuf, prim["attributes"]["POSITION"]).astype(np.float64)
            pos = pos @ world[:3, :3].T + world[:3, 3]
            # glTF Y-up -> Blender Z-up: (x, y, z) -> (x, -z, y)
            pos = np.column_stack([pos[:, 0], -pos[:, 2], pos[:, 1]])
            if "indices" in prim:
                tris = read_accessor(gltf, binbuf, prim["indices"]).astype(np.int64).reshape(-1, 3)
            else:
                tris = np.arange(len(pos), dtype=np.int64).reshape(-1, 3)
            uv_acc = prim["attributes"].get("TEXCOORD_0")
            if uv_acc is not None:
                uvs = read_accessor(gltf, binbuf, uv_acc).astype(np.float64)
            else:
                uvs = np.zeros((len(pos), 2))
            verts_list.append(pos)
            tris_list.append(tris + v_offset)
            mat_list.append(np.full(len(tris), prim.get("material", -1), dtype=np.int64))
            uv_list.append(uvs)
            v_offset += len(pos)

    if not verts_list:
        return None

    V = np.vstack(verts_list)
    F = np.vstack(tris_list)
    M = np.concatenate(mat_list)
    UV = np.vstack(uv_list)

    # Remap glTF material indices to compact slot indices.
    used = [int(i) for i in np.unique(M)]
    remap = {g: s for s, g in enumerate(used)}
    M_slots = np.vectorize(remap.get)(M) if len(used) > 1 else np.zeros(len(M), dtype=np.int64)

    # Per-loop UVs with glTF->Blender V flip.
    loop_uvs = UV[F.ravel()]
    loop_uvs[:, 1] = 1.0 - loop_uvs[:, 1]

    return V, F, M_slots, loop_uvs, used


# ----------------------------
# Blender-side building
# ----------------------------

def get_base_color(gltf, mat):
    """Return (rgba list, texture image index or None) for a glTF material."""
    pbr = mat.get("pbrMetallicRoughness", {})
    color = pbr.get("baseColorFactor", [0.8, 0.8, 0.8, 1.0])
    tex_index = None
    tex = pbr.get("baseColorTexture")
    # SimLab sometimes writes spec/gloss materials.
    sg = mat.get("extensions", {}).get("KHR_materials_pbrSpecularGlossiness")
    if sg:
        color = sg.get("diffuseFactor", color)
        tex = sg.get("diffuseTexture", tex)
    if tex is not None:
        tex_index = gltf["textures"][tex["index"]].get("source")
    return color, tex_index


def load_embedded_image(gltf, binbuf, image_index, cache):
    """Extract an embedded image, load into Blender, pack, return bpy image."""
    if image_index in cache:
        return cache[image_index]
    img_def = gltf["images"][image_index]
    if "bufferView" not in img_def:
        return None  # external URI images not expected from these exporters
    existing = bpy.data.images.get(img_def.get("name") or f"skp_tex_{image_index}")
    if existing is not None:
        cache[image_index] = existing
        return existing
    bv = gltf["bufferViews"][img_def["bufferView"]]
    start = bv.get("byteOffset", 0)
    raw = binbuf[start:start + bv["byteLength"]]
    ext = ".png" if img_def.get("mimeType") == "image/png" else ".jpg"
    name = img_def.get("name") or f"skp_tex_{image_index}"
    tmp = tempfile.NamedTemporaryFile(suffix=ext, delete=False)
    try:
        tmp.write(raw)
        tmp.close()
        image = bpy.data.images.load(tmp.name)
        image.name = name
        image.pack()
    finally:
        try:
            os.unlink(tmp.name)
        except OSError:
            pass
    cache[image_index] = image
    return image


def get_or_create_material(gltf, binbuf, gltf_mat_index, mat_cache, img_cache):
    """Build (or reuse) a Blender material for a glTF material index."""
    if gltf_mat_index < 0:
        key = "__skp_default__"
        if key not in mat_cache:
            mat = bpy.data.materials.new("SKP Default")
            mat.use_nodes = True
            mat_cache[key] = mat
        return mat_cache[key]

    mat_def = gltf["materials"][gltf_mat_index]
    name = mat_def.get("name", f"SKP Material {gltf_mat_index}")
    if name in mat_cache:
        return mat_cache[name]

    # Cross-run reuse: a material of this name from an earlier import
    # (e.g. another building in the same .blend) is the same SketchUp
    # material. Reuse it instead of creating "Name.001" duplicates.
    existing = bpy.data.materials.get(name)
    if existing is not None:
        mat_cache[name] = existing
        return existing

    mat = bpy.data.materials.new(name)
    mat.use_nodes = True
    bsdf = next(n for n in mat.node_tree.nodes if n.type == 'BSDF_PRINCIPLED')
    color, tex_index = get_base_color(gltf, mat_def)
    bsdf.inputs["Base Color"].default_value = color
    mat.diffuse_color = color  # solid-mode viewport color
    if len(color) > 3 and color[3] < 1.0:
        bsdf.inputs["Alpha"].default_value = color[3]

    if tex_index is not None:
        image = load_embedded_image(gltf, binbuf, tex_index, img_cache)
        if image:
            tex_node = mat.node_tree.nodes.new('ShaderNodeTexImage')
            tex_node.image = image
            tex_node.location = (-320, 260)
            mat.node_tree.links.new(tex_node.outputs["Color"],
                                    bsdf.inputs["Base Color"])

    mat_cache[name] = mat
    return mat


def build_mesh_object(name, geo, gltf, binbuf, mat_cache, img_cache):
    V, F, M_slots, loop_uvs, used_mats = geo
    n_tris = len(F)

    mesh = bpy.data.meshes.new(name)
    mesh.vertices.add(len(V))
    mesh.vertices.foreach_set("co", V.astype(np.float32).ravel())
    mesh.loops.add(n_tris * 3)
    mesh.loops.foreach_set("vertex_index", F.astype(np.int32).ravel())
    mesh.polygons.add(n_tris)
    mesh.polygons.foreach_set("loop_start",
                              np.arange(0, n_tris * 3, 3, dtype=np.int32))
    try:
        mesh.polygons.foreach_set("loop_total",
                                  np.full(n_tris, 3, dtype=np.int32))
    except (AttributeError, RuntimeError):
        pass  # read-only/derived in newer Blender versions
    mesh.polygons.foreach_set("material_index", M_slots.astype(np.int32))

    uv_layer = mesh.uv_layers.new(name="UVMap")
    uv_layer.data.foreach_set("uv", loop_uvs.astype(np.float32).ravel())

    mesh.update(calc_edges=True)
    mesh.validate(verbose=False)

    for gltf_mat_index in used_mats:
        mesh.materials.append(
            get_or_create_material(gltf, binbuf, gltf_mat_index,
                                   mat_cache, img_cache))

    obj = bpy.data.objects.new(name, mesh)
    return obj


def fallback_stock_import(context, filepath, stem, target_coll):
    """Slow path: stock glTF importer + collapse. Used only if parsing fails."""
    before = set(bpy.data.objects)
    bpy.ops.import_scene.gltf(filepath=filepath)
    new_objs = [o for o in bpy.data.objects if o not in before]
    mesh_objs = [o for o in new_objs if o.type == 'MESH']
    others = [o for o in new_objs if o.type != 'MESH']

    joined = None
    if mesh_objs:
        for obj in context.selected_objects:
            obj.select_set(False)
        for obj in mesh_objs:
            obj.select_set(True)
        context.view_layer.objects.active = mesh_objs[0]
        if len(mesh_objs) > 1:
            bpy.ops.object.join()
        joined = mesh_objs[0]
        joined.select_set(False)
        joined.name = stem
        if joined.parent:
            world = joined.matrix_world.copy()
            joined.parent = None
            joined.matrix_world = world
        for coll in list(joined.users_collection):
            coll.objects.unlink(joined)
        target_coll.objects.link(joined)

    if others:
        bpy.data.batch_remove(others)
    return joined


# ----------------------------
# Properties / Operator / UI
# ----------------------------

class SKPGltfImportSettings(PropertyGroup):
    import_dir: StringProperty(
        name="Folder",
        description="Folder containing per-tag .glb files",
        subtype='DIR_PATH',
        default="",
    )
    scale_factor: FloatProperty(
        name="Scale Factor (deprecated)",
        description=(
            "DEPRECATED: DAE-era scale fix. glTF carries real units; "
            "leave at 1.0"
        ),
        default=1.0,
        min=0.0001,
        max=10000.0,
    )


class SKP_OT_import_tag_gltf(Operator):
    bl_idname = "skp.import_tag_gltf"
    bl_label = "Import Tag .glb Files"
    bl_description = "Import every .glb in the folder, one object per file"
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return context.mode == 'OBJECT'

    def execute(self, context):
        settings = context.scene.skp_gltf_import
        folder = bpy.path.abspath(settings.import_dir)

        if not folder or not os.path.isdir(folder):
            self.report({'ERROR'}, "Select a valid folder first.")
            return {'CANCELLED'}

        files = sorted(
            f for f in os.listdir(folder)
            if f.lower().endswith(GLTF_EXTENSIONS)
        )
        if not files:
            self.report({'ERROR'}, "No .glb files found in that folder.")
            return {'CANCELLED'}

        stamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
        run_coll = bpy.data.collections.new(f"SKP Import {stamp}")
        context.scene.collection.children.link(run_coll)

        mat_cache, img_cache = {}, {}
        imported, failed = 0, []
        t0 = datetime.datetime.now()

        for filename in files:
            filepath = os.path.join(folder, filename)
            stem = os.path.splitext(filename)[0]
            try:
                gltf, binbuf = parse_glb(filepath)
                geo = collect_geometry(gltf, binbuf)
                if geo is None:
                    print(f"[SKP Import] No mesh geometry in {filename}")
                    continue
                obj = build_mesh_object(stem, geo, gltf, binbuf,
                                        mat_cache, img_cache)
                if settings.scale_factor != 1.0:
                    obj.scale = (settings.scale_factor,) * 3
                run_coll.objects.link(obj)
                imported += 1
                print(f"[SKP Import] {filename}: "
                      f"{len(geo[1]):,} tris, {len(geo[4])} materials")
            except Exception as e:
                print(f"[SKP Import] Direct parse failed for {filename} ({e}), "
                      f"trying stock importer...")
                try:
                    if fallback_stock_import(context, filepath, stem, run_coll):
                        imported += 1
                except Exception as e2:
                    failed.append(filename)
                    print(f"[SKP Import] FAILED {filename}: {e2}")

        elapsed = (datetime.datetime.now() - t0).total_seconds()
        msg = (f"Imported {imported} of {len(files)} files "
               f"in {elapsed:.1f}s into '{run_coll.name}'.")
        if failed:
            msg += f" Failed: {', '.join(failed)} (see console)."
            self.report({'WARNING'}, msg)
        else:
            self.report({'INFO'}, msg)
        return {'FINISHED'}


class SKP_PT_import_panel(Panel):
    bl_label = "SketchUp to Blender"
    bl_idname = "SKP_PT_import_panel"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "SKP Import"

    def draw(self, context):
        layout = self.layout
        settings = context.scene.skp_gltf_import
        layout.prop(settings, "import_dir")
        layout.prop(settings, "scale_factor")
        layout.operator(SKP_OT_import_tag_gltf.bl_idname, icon='IMPORT')


classes = (
    SKPGltfImportSettings,
    SKP_OT_import_tag_gltf,
    SKP_PT_import_panel,
)


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    bpy.types.Scene.skp_gltf_import = bpy.props.PointerProperty(type=SKPGltfImportSettings)


def unregister():
    del bpy.types.Scene.skp_gltf_import
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)


if __name__ == "__main__":
    register()