import os
import struct
from collections import defaultdict
from typing import Optional, Tuple

import numpy as np
import trimesh
from PIL import Image

# =========================
# 你只需要改这里
# =========================
# INPUT_GLB = "/mnt/pfs/users/huangzehuan/projects/SegviGen/examples/trellis2_output.glb"


# UID = "demonic_warrior_red_bronze_armor"
# UID = "playful_pose_white_top_portrait"
# UID = "african_inspired_metallic_silver_ensemble_with_headwrap"
# UID = "cyberpunk_bowser_motorcycle"
# UID = "crimson_battle_mecha_with_spikes"
UID = "black_lace_lingerie_ensemble"

INPUT_GLB = (
    f"/mnt/pfs/users/maxueqi/studio/datasets/dense_mesh/segvigen_bak/{UID}/output.glb"
)

# 只用 RGB（忽略透明度/alpha）
COLOR_QUANT_STEP = 16  # RGB 量化步长：0/4/8/16（越大越“合并”）
PALETTE_SAMPLE_PIXELS = 2_000_000
PALETTE_MIN_PIXELS = 500  # 少于该像素数的颜色当噪声丢掉（边界抗锯齿中间色）
PALETTE_MAX_COLORS = 256  # 最多保留多少个主颜色
PALETTE_MERGE_DIST = 32  # ✅ 合并 palette 内近似颜色（解决“看着同色却拆两块”）

SAMPLES_PER_FACE = 4  # 1 或 4（推荐 4，能明显减少边界采样误差）
FLIP_V = True  # glTF 常见需要 flip V
UV_WRAP_REPEAT = True  # True: repeat (mod 1)；False: clamp 到 [0,1]

MIN_FACES_PER_PART = 50
BAKE_TRANSFORMS = True
DEBUG_PRINT = True
# =========================


CHUNK_TYPE_JSON = 0x4E4F534A  # b'JSON'
CHUNK_TYPE_BIN = 0x004E4942  # b'BIN\0'


def _default_out_path(in_path: str) -> str:
    root, ext = os.path.splitext(in_path)
    if ext.lower() not in [".glb", ".gltf"]:
        ext = ".glb"
    return f"{root}_seg.glb"


def _quantize_rgb(rgb: np.ndarray, step: int) -> np.ndarray:
    """
    rgb: (...,3) uint8
    """
    if step is None or step <= 0:
        return rgb
    q = (rgb.astype(np.int32) + step // 2) // step * step
    q = np.clip(q, 0, 255).astype(np.uint8)
    return q


def _load_glb_json_and_bin(glb_path: str) -> Tuple[dict, bytes]:
    data = open(glb_path, "rb").read()
    if len(data) < 12:
        raise RuntimeError("Invalid GLB: too small")

    magic, version, length = struct.unpack_from("<4sII", data, 0)
    if magic != b"glTF":
        raise RuntimeError("Not a GLB file (missing glTF header)")

    offset = 12
    gltf_json = None
    bin_chunk = None

    while offset + 8 <= len(data):
        chunk_len, chunk_type = struct.unpack_from("<II", data, offset)
        offset += 8
        chunk_data = data[offset : offset + chunk_len]
        offset += chunk_len

        if chunk_type == CHUNK_TYPE_JSON:
            gltf_json = chunk_data.decode("utf-8", errors="replace")
        elif chunk_type == CHUNK_TYPE_BIN:
            bin_chunk = chunk_data

    if gltf_json is None:
        raise RuntimeError("GLB missing JSON chunk")
    if bin_chunk is None:
        raise RuntimeError("GLB missing BIN chunk")

    import json

    return json.loads(gltf_json), bin_chunk


def _extract_basecolor_texture_image(glb_path: str) -> np.ndarray:
    """
    从 GLB 内嵌资源里拿 baseColorTexture 的 PNG/JPG，返回 (H,W,4) uint8 RGBA
    """
    gltf, bin_chunk = _load_glb_json_and_bin(glb_path)

    materials = gltf.get("materials", [])
    textures = gltf.get("textures", [])
    images = gltf.get("images", [])
    buffer_views = gltf.get("bufferViews", [])

    if not materials:
        raise RuntimeError("No materials in GLB")

    # 这里按 material[0] 取 baseColorTexture（你的 glb 只有一个材质/primitive）
    pbr = materials[0].get("pbrMetallicRoughness", {})
    base_tex_index = pbr.get("baseColorTexture", {}).get("index", None)
    if base_tex_index is None:
        raise RuntimeError("Material has no baseColorTexture")

    if base_tex_index >= len(textures):
        raise RuntimeError("baseColorTexture index out of range")

    tex = textures[base_tex_index]
    img_index = tex.get("source", None)
    if img_index is None or img_index >= len(images):
        raise RuntimeError("Texture has no valid image source")

    img_info = images[img_index]
    bv_index = img_info.get("bufferView", None)
    mime = img_info.get("mimeType", None)
    if bv_index is None:
        uri = img_info.get("uri", None)
        raise RuntimeError(f"Image is not embedded (bufferView missing). uri={uri}")

    if bv_index >= len(buffer_views):
        raise RuntimeError("image.bufferView out of range")

    bv = buffer_views[bv_index]
    bo = int(bv.get("byteOffset", 0))
    bl = int(bv.get("byteLength", 0))
    img_bytes = bin_chunk[bo : bo + bl]

    if DEBUG_PRINT:
        print(
            f"[Texture] baseColorTextureIndex={base_tex_index}, imageIndex={img_index}, "
            f"bufferView={bv_index}, mime={mime}, bytes={len(img_bytes)}"
        )

    pil = Image.open(trimesh.util.wrap_as_stream(img_bytes)).convert("RGBA")
    return np.array(pil, dtype=np.uint8)


def _merge_palette_rgb(
    palette_rgb: np.ndarray, counts: np.ndarray, merge_dist: float
) -> np.ndarray:
    """
    对 palette 内 RGB 做“近似合并”，用 counts 作为权重更新中心。
    palette_rgb: (K,3) uint8
    counts: (K,) int
    """
    if palette_rgb is None or len(palette_rgb) == 0:
        return palette_rgb
    if merge_dist is None or merge_dist <= 0:
        return palette_rgb

    rgb = palette_rgb.astype(np.float32)
    counts = counts.astype(np.int64)

    order = np.argsort(-counts)

    centers = []
    center_w = []
    thr2 = float(merge_dist) * float(merge_dist)

    for idx in order:
        x = rgb[idx]
        w = int(counts[idx])

        if not centers:
            centers.append(x.copy())
            center_w.append(w)
            continue

        C = np.stack(centers, axis=0)  # (M,3)
        d2 = np.sum((C - x[None, :]) ** 2, axis=1)
        k = int(np.argmin(d2))

        if float(d2[k]) <= thr2:
            cw = center_w[k]
            centers[k] = (centers[k] * cw + x * w) / (cw + w)
            center_w[k] = cw + w
        else:
            centers.append(x.copy())
            center_w.append(w)

    merged = np.clip(np.rint(np.stack(centers, axis=0)), 0, 255).astype(np.uint8)

    if DEBUG_PRINT:
        print(
            f"[PaletteMerge] before={len(palette_rgb)} after={len(merged)} merge_dist={merge_dist}"
        )

    return merged


def _build_palette_rgb(tex_rgba: np.ndarray) -> np.ndarray:
    """
    从贴图中提取 RGB 主颜色调色板（忽略 alpha）。
    返回: (K,3) uint8
    """
    rgb = tex_rgba[:, :, :3].reshape(-1, 3)
    n = rgb.shape[0]

    if n > PALETTE_SAMPLE_PIXELS:
        rng = np.random.default_rng(0)
        idx = rng.choice(n, size=PALETTE_SAMPLE_PIXELS, replace=False)
        rgb = rgb[idx]

    rgb = _quantize_rgb(rgb, COLOR_QUANT_STEP)

    uniq, counts = np.unique(rgb, axis=0, return_counts=True)
    order = np.argsort(-counts)
    uniq = uniq[order]
    counts = counts[order]

    keep = counts >= PALETTE_MIN_PIXELS
    uniq = uniq[keep]
    counts = counts[keep]

    if len(uniq) > PALETTE_MAX_COLORS:
        uniq = uniq[:PALETTE_MAX_COLORS]
        counts = counts[:PALETTE_MAX_COLORS]

    if DEBUG_PRINT:
        print(
            f"[Palette] quant_step={COLOR_QUANT_STEP} palette_size(before_merge)={len(uniq)} "
            f"min_pixels={PALETTE_MIN_PIXELS}"
        )
        for i in range(min(15, len(uniq))):
            r, g, b = [int(x) for x in uniq[i]]
            print(f"  {i:02d} rgb=({r},{g},{b}) count={int(counts[i])}")

    uniq = _merge_palette_rgb(uniq.astype(np.uint8), counts, PALETTE_MERGE_DIST)

    if DEBUG_PRINT:
        print(f"[Palette] palette_size(after_merge)={len(uniq)}")
        for i in range(min(15, len(uniq))):
            r, g, b = [int(x) for x in uniq[i]]
            print(f"  {i:02d} rgb=({r},{g},{b})")

    return uniq.astype(np.uint8)


def _unwrap_uv3_for_seam(uv3: np.ndarray) -> np.ndarray:
    """
    uv3: (F,3,2). 若跨 seam（跨度>0.5），把小于0.5的一侧 +1，避免均值跑到另一边。
    """
    out = uv3.copy()
    for d in range(2):
        v = out[:, :, d]
        vmin = v.min(axis=1)
        vmax = v.max(axis=1)
        seam = (vmax - vmin) > 0.5
        if np.any(seam):
            vv = v[seam]
            vv = np.where(vv < 0.5, vv + 1.0, vv)
            out[seam, :, d] = vv
    return out


def _barycentric_samples(uv3: np.ndarray, samples_per_face: int) -> np.ndarray:
    """
    uv3: (F,3,2)
    return: (F,S,2)
    """
    uv3 = _unwrap_uv3_for_seam(uv3)

    if samples_per_face == 1:
        w = np.array([1 / 3, 1 / 3, 1 / 3], dtype=np.float32)
        uvs = uv3[:, 0, :] * w[0] + uv3[:, 1, :] * w[1] + uv3[:, 2, :] * w[2]
        return uvs[:, None, :]

    # 4 个点：中心 + 三个靠近顶点的内点（尽量远离边界抗锯齿带）
    ws = np.array(
        [
            [1 / 3, 1 / 3, 1 / 3],
            [0.80, 0.10, 0.10],
            [0.10, 0.80, 0.10],
            [0.10, 0.10, 0.80],
        ],
        dtype=np.float32,
    )
    uvs = (
        uv3[:, None, 0, :] * ws[None, :, 0, None]
        + uv3[:, None, 1, :] * ws[None, :, 1, None]
        + uv3[:, None, 2, :] * ws[None, :, 2, None]
    )
    return uvs


def _wrap_or_clamp_uv(uv: np.ndarray) -> np.ndarray:
    if UV_WRAP_REPEAT:
        return np.mod(uv, 1.0)
    return np.clip(uv, 0.0, 1.0)


def _sample_texture_nearest_rgb(tex_rgba: np.ndarray, uv: np.ndarray) -> np.ndarray:
    """
    tex_rgba: (H,W,4) uint8
    uv: (N,2) float
    return: (N,3) uint8
    """
    h, w = tex_rgba.shape[0], tex_rgba.shape[1]
    uv = _wrap_or_clamp_uv(uv)

    u = uv[:, 0]
    v = uv[:, 1]
    if FLIP_V:
        v = 1.0 - v

    x = np.rint(u * (w - 1)).astype(np.int32)
    y = np.rint(v * (h - 1)).astype(np.int32)
    x = np.clip(x, 0, w - 1)
    y = np.clip(y, 0, h - 1)

    return tex_rgba[y, x, :3].astype(np.uint8)


def _map_to_palette_rgb(
    colors_rgb: np.ndarray, palette_rgb: np.ndarray, chunk: int = 20000
) -> Tuple[np.ndarray, np.ndarray]:
    """
    把采样到的 RGB 映射到最近的 palette RGB.
    如果 palette 为空，则用 colors_rgb 的 unique 作为“临时 palette”.
    返回:
      labels: (N,) int
      used_palette_rgb: (K,3) uint8
    """
    if palette_rgb is None or len(palette_rgb) == 0:
        uniq, inv = np.unique(colors_rgb, axis=0, return_inverse=True)
        return inv.astype(np.int32), uniq.astype(np.uint8)

    c = colors_rgb.astype(np.float32)
    p = palette_rgb.astype(np.float32)

    out = np.empty((c.shape[0],), dtype=np.int32)
    for i in range(0, c.shape[0], chunk):
        cc = c[i : i + chunk]
        d2 = ((cc[:, None, :] - p[None, :, :]) ** 2).sum(axis=2)
        out[i : i + chunk] = np.argmin(d2, axis=1).astype(np.int32)

    return out, palette_rgb


def _face_labels_from_texture_rgb(
    mesh: trimesh.Trimesh,
    tex_rgba: np.ndarray,
    palette_rgb: np.ndarray,
) -> Optional[Tuple[np.ndarray, np.ndarray]]:
    """
    用 TEXCOORD_0 + baseColorTexture，为每个 face 采样 RGB，并映射到 palette label。
    返回:
      face_label: (F,) int
      label_rgb: (K,3) uint8
    """
    uv = getattr(mesh.visual, "uv", None)
    if uv is None:
        return None

    uv = np.asarray(uv, dtype=np.float32)
    if uv.ndim != 2 or uv.shape[1] != 2 or uv.shape[0] != len(mesh.vertices):
        return None

    faces = mesh.faces
    uv3 = uv[faces]  # (F,3,2)

    uvs = _barycentric_samples(uv3, SAMPLES_PER_FACE)  # (F,S,2)
    F, S = uvs.shape[0], uvs.shape[1]
    flat_uv = uvs.reshape(-1, 2)

    sampled_rgb = _sample_texture_nearest_rgb(tex_rgba, flat_uv)  # (F*S,3)
    sampled_rgb = _quantize_rgb(sampled_rgb, COLOR_QUANT_STEP)

    sample_label, used_palette = _map_to_palette_rgb(sampled_rgb, palette_rgb)
    sample_label = sample_label.reshape(F, S)

    if S == 1:
        return sample_label[:, 0].astype(np.int32), used_palette

    # 4 票投票（向量化）
    l0, l1, l2, l3 = (
        sample_label[:, 0],
        sample_label[:, 1],
        sample_label[:, 2],
        sample_label[:, 3],
    )
    c0 = 1 + (l0 == l1) + (l0 == l2) + (l0 == l3)
    c1 = 1 + (l1 == l0) + (l1 == l2) + (l1 == l3)
    c2 = 1 + (l2 == l0) + (l2 == l1) + (l2 == l3)
    c3 = 1 + (l3 == l0) + (l3 == l1) + (l3 == l2)

    counts = np.stack([c0, c1, c2, c3], axis=1)  # (F,4)
    vals = np.stack([l0, l1, l2, l3], axis=1)  # (F,4)
    best = vals[np.arange(F), np.argmax(counts, axis=1)]
    return best.astype(np.int32), used_palette


# =========================
# 拓扑纠错
# =========================

import numpy as np
import trimesh
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components


def _get_physical_face_adjacency(mesh: trimesh.Trimesh) -> np.ndarray:
    """
    忽略 UV 接缝，计算纯物理空间上的面片相邻关系。
    """
    # 1. 四舍五入顶点坐标（处理浮点数微小误差），找出空间中真正唯一的物理顶点
    v_rounded = np.round(mesh.vertices, decimals=3)
    v_unique, inv_indices = np.unique(v_rounded, axis=0, return_inverse=True)

    # 2. 将原本的面片索引，映射到这些“唯一物理顶点”上
    # 这样，跨越 UV 接缝的面片，此时它们引用的顶点索引就变成一样的了
    physical_faces = inv_indices[mesh.faces]

    # 3. 创建一个临时的“影子网格”（process=False 极其重要，防止 trimesh 内部重排面片）
    tmp_mesh = trimesh.Trimesh(vertices=v_unique, faces=physical_faces, process=False)

    # 返回影子网格的物理相邻边
    return tmp_mesh.face_adjacency


def smooth_face_labels_by_topology(
    mesh: trimesh.Trimesh, face_label: np.ndarray, min_faces: int = 50
) -> np.ndarray:
    """
    通过真实的 3D 物理拓扑关系过滤飞点，跨越 UV 接缝合并色块。

    Phase 1: 在同色连通图上，把挨着大块的小块吞并到大块中。
    Phase 2: 对残留小块（邻居全是小块），回退到全物理邻接，
             按物理邻居中的多数 label 吞并。
    Phase 3: 对完全孤立的面片（无物理邻接边），按面片质心距离
             找最近的非孤立面片，继承其 label。
    """
    labels = face_label.copy()
    edges = _get_physical_face_adjacency(mesh)
    F = len(mesh.faces)

    # ---- Phase 1: 同色连通域平滑 ----
    for iteration in range(3):
        same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
        sub_edges = edges[same_label]

        if len(sub_edges) > 0:
            data = np.ones(len(sub_edges), dtype=bool)
            graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
            graph = graph.maximum(graph.T)
            n_components, comp_labels = connected_components(graph, directed=False)
        else:
            n_components = F
            comp_labels = np.arange(F)

        comp_sizes = np.bincount(comp_labels, minlength=n_components)
        small_comps = np.where(comp_sizes < min_faces)[0]
        if len(small_comps) == 0:
            break

        is_small = np.isin(comp_labels, small_comps)

        mask0 = is_small[edges[:, 0]]
        mask1 = is_small[edges[:, 1]]

        boundary_edges_0 = edges[mask0 & ~mask1]
        boundary_edges_1 = edges[mask1 & ~mask0]

        b_inner = np.concatenate([boundary_edges_0[:, 0], boundary_edges_1[:, 1]])
        b_outer = np.concatenate([boundary_edges_0[:, 1], boundary_edges_1[:, 0]])

        if len(b_inner) == 0:
            break

        outer_labels = labels[b_outer]
        inner_comps = comp_labels[b_inner]

        for cid in np.unique(inner_comps):
            cid_mask = inner_comps == cid
            surrounding_labels = outer_labels[cid_mask]
            if len(surrounding_labels) > 0:
                best_label = np.bincount(surrounding_labels).argmax()
                labels[comp_labels == cid] = best_label

    # ---- Phase 2: 用全物理邻接处理残留小块 ----
    # 重新计算同色连通域，找出还残留的小块
    same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
    sub_edges = edges[same_label]
    if len(sub_edges) > 0:
        data = np.ones(len(sub_edges), dtype=bool)
        graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
        graph = graph.maximum(graph.T)
        n_components, comp_labels = connected_components(graph, directed=False)
    else:
        n_components = F
        comp_labels = np.arange(F)

    comp_sizes = np.bincount(comp_labels, minlength=n_components)
    small_comps_set = set(np.where(comp_sizes < min_faces)[0])

    if small_comps_set:
        is_small = np.array([comp_labels[i] in small_comps_set for i in range(F)])

        # 构建全物理邻接查找表: face -> set of neighbor faces
        adj = defaultdict(set)
        for e0, e1 in edges:
            adj[int(e0)].add(int(e1))
            adj[int(e1)].add(int(e0))

        # 迭代：每轮让小块面片从物理邻居（忽略颜色）中投票取多数 label
        for _ in range(3):
            changed = False
            small_comps_now = set(
                int(c)
                for c in range(n_components)
                if comp_sizes[c] < min_faces and c in small_comps_set
            )
            if not small_comps_now:
                break

            for cid in small_comps_now:
                cid_faces = np.where(comp_labels == cid)[0]
                # 收集所有物理邻居中不属于本连通域的面片的 label
                neighbor_labels = []
                for fi in cid_faces:
                    for nf in adj[int(fi)]:
                        if comp_labels[nf] != cid:
                            neighbor_labels.append(labels[nf])

                if len(neighbor_labels) > 0:
                    best_label = int(np.bincount(neighbor_labels).argmax())
                    labels[cid_faces] = best_label
                    changed = True

            if not changed:
                break

            # 重新计算连通域
            same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
            sub_edges = edges[same_label]
            if len(sub_edges) > 0:
                data = np.ones(len(sub_edges), dtype=bool)
                graph = coo_matrix(
                    (data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F)
                )
                graph = graph.maximum(graph.T)
                n_components, comp_labels = connected_components(graph, directed=False)
            else:
                n_components = F
                comp_labels = np.arange(F)
            comp_sizes = np.bincount(comp_labels, minlength=n_components)
            small_comps_set = set(np.where(comp_sizes < min_faces)[0])

    # ---- Phase 3: 完全孤立面片（无物理邻接边），按质心距离继承 label ----
    same_label = labels[edges[:, 0]] == labels[edges[:, 1]]
    sub_edges = edges[same_label]
    if len(sub_edges) > 0:
        data = np.ones(len(sub_edges), dtype=bool)
        graph = coo_matrix((data, (sub_edges[:, 0], sub_edges[:, 1])), shape=(F, F))
        graph = graph.maximum(graph.T)
        _, comp_labels = connected_components(graph, directed=False)
    else:
        comp_labels = np.arange(F)
    comp_sizes = np.bincount(comp_labels)
    orphan_comps = set(np.where(comp_sizes < min_faces)[0])

    if orphan_comps:
        orphan_mask = np.array([comp_labels[i] in orphan_comps for i in range(F)])
        non_orphan_mask = ~orphan_mask
        if non_orphan_mask.any() and orphan_mask.any():
            centroids = mesh.triangles_center
            orphan_indices = np.where(orphan_mask)[0]
            non_orphan_indices = np.where(non_orphan_mask)[0]
            non_orphan_centroids = centroids[non_orphan_indices]

            for oi in orphan_indices:
                dists = np.linalg.norm(non_orphan_centroids - centroids[oi], axis=1)
                nearest = non_orphan_indices[np.argmin(dists)]
                labels[oi] = labels[nearest]

        if DEBUG_PRINT:
            n_orphan = int(orphan_mask.sum())
            print(f"  [Phase3] Assigned {n_orphan} orphan faces by centroid proximity")

    return labels


# =========================
# 分割主函数
# =========================


# def split_glb_by_texture_palette_rgb(
#     in_glb_path: str,
#     out_glb_path: Optional[str] = None,
#     min_faces_per_part: int = 1,
#     bake_transforms: bool = True,
# ) -> str:
#     """
#     输入：glb（无 COLOR_0，但有 baseColorTexture + TEXCOORD_0）
#     输出：先从贴图提取 RGB 主色 palette（忽略 alpha），再按 palette label 分割
#     """
#     if out_glb_path is None:
#         out_glb_path = _default_out_path(in_glb_path)

#     tex_rgba = _extract_basecolor_texture_image(in_glb_path)
#     palette_rgb = _build_palette_rgb(tex_rgba)

#     scene = trimesh.load(in_glb_path, force="scene", process=False)
#     out_scene = trimesh.Scene()

#     part_count = 0
#     base = os.path.splitext(os.path.basename(in_glb_path))[0]

#     for node_name in scene.graph.nodes_geometry:
#         geom_name = scene.graph[node_name][1]
#         if geom_name is None:
#             continue

#         geom = scene.geometry.get(geom_name, None)
#         if geom is None or not isinstance(geom, trimesh.Trimesh):
#             continue

#         mesh = geom.copy()

#         if bake_transforms:
#             T, _ = scene.graph.get(node_name)
#             if T is not None:
#                 mesh.apply_transform(T)

#         res = _face_labels_from_texture_rgb(mesh, tex_rgba, palette_rgb)
#         if res is None:
#             if DEBUG_PRINT:
#                 print(f"[{node_name}] no uv / cannot sample -> keep orig")
#             out_scene.add_geometry(mesh, geom_name=f"{base}__{node_name}__orig")
#             continue

#         face_label, label_rgb = res

#         # =========================
#         # 🔥 新增调用：进行拓扑纠错，合并飞点
#         # =========================
#         face_label = smooth_face_labels_by_topology(mesh, face_label, min_faces=100)

#         if DEBUG_PRINT:
#             uniq_labels, cnts = np.unique(face_label, return_counts=True)
#             order = np.argsort(-cnts)
#             print(
#                 f"[{node_name}] faces={len(mesh.faces)} labels_used={len(uniq_labels)} palette_size={len(label_rgb)}"
#             )
#             for i in order[:10]:
#                 lab = int(uniq_labels[i])
#                 r, g, b = (
#                     [int(x) for x in label_rgb[lab]]
#                     if 0 <= lab < len(label_rgb)
#                     else (0, 0, 0)
#                 )
#                 print(f"  label={lab} rgb=({r},{g},{b}) faces={int(cnts[i])}")

#         groups = defaultdict(list)
#         for fi, lab in enumerate(face_label):
#             groups[int(lab)].append(fi)

#         for lab, face_ids in groups.items():
#             if len(face_ids) < min_faces_per_part:
#                 continue

#             sub = mesh.submesh(
#                 [np.array(face_ids, dtype=np.int64)], append=True, repair=False
#             )
#             if sub is None:
#                 continue
#             if isinstance(sub, (list, tuple)):
#                 if not sub:
#                     continue
#                 sub = sub[0]

#             if 0 <= lab < len(label_rgb):
#                 r, g, b = [int(x) for x in label_rgb[lab]]
#                 part_name = f"{base}__{node_name}__label_{lab}__rgb_{r}_{g}_{b}"
#             else:
#                 part_name = f"{base}__{node_name}__label_{lab}"

#             out_scene.add_geometry(sub, geom_name=part_name)
#             part_count += 1

#     if part_count == 0:
#         if DEBUG_PRINT:
#             print("[INFO] part_count==0, fallback to original scene export.")
#         out_scene = scene

#     out_scene.export(out_glb_path)
#     return out_glb_path


def split_glb_by_texture_palette_rgb(
    in_glb_path: str,
    out_glb_path: Optional[str] = None,
    min_faces_per_part: int = 1,
    bake_transforms: bool = True,
    color_quant_step: int = 16,
    palette_sample_pixels: int = 2_000_000,
    palette_min_pixels: int = 500,
    palette_max_colors: int = 256,
    palette_merge_dist: int = 32,
    samples_per_face: int = 4,
    flip_v: bool = True,
    uv_wrap_repeat: bool = True,
    transition_conf_thresh: float = 1.0,
    transition_prop_iters: int = 6,
    transition_neighbor_min: int = 1,
    small_component_action: str = "reassign",
    small_component_min_faces: int = 50,
    postprocess_iters: int = 3,
    debug_print: bool = True,
) -> str:
    """
    Input: GLB (no COLOR_0, but with baseColorTexture + TEXCOORD_0)
    Output: Split based on palette labels derived from baseColorTexture
    """
    if out_glb_path is None:
        out_glb_path = _default_out_path(in_glb_path)

    tex_rgba = _extract_basecolor_texture_image(in_glb_path)
    palette_rgb = _build_palette_rgb(tex_rgba)

    scene = trimesh.load(in_glb_path, force="scene", process=False)
    out_scene = trimesh.Scene()

    part_count = 0
    base = os.path.splitext(os.path.basename(in_glb_path))[0]

    for node_name in scene.graph.nodes_geometry:
        geom_name = scene.graph[node_name][1]
        if geom_name is None:
            continue

        geom = scene.geometry.get(geom_name, None)
        if geom is None or not isinstance(geom, trimesh.Trimesh):
            continue

        mesh = geom.copy()

        if bake_transforms:
            T, _ = scene.graph.get(node_name)
            if T is not None:
                mesh.apply_transform(T)

        res = _face_labels_from_texture_rgb(mesh, tex_rgba, palette_rgb)
        if res is None:
            if debug_print:
                print(f"[{node_name}] no uv / cannot sample -> keep orig")
            out_scene.add_geometry(mesh, geom_name=f"{base}__{node_name}__orig")
            continue

        face_label, label_rgb = res

        # =========================
        # 🔥 New: Apply topology correction to merge small disconnected components
        # =========================
        face_label = smooth_face_labels_by_topology(mesh, face_label, min_faces=100)

        if debug_print:
            uniq_labels, cnts = np.unique(face_label, return_counts=True)
            order = np.argsort(-cnts)
            print(
                f"[{node_name}] faces={len(mesh.faces)} labels_used={len(uniq_labels)} palette_size={len(label_rgb)}"
            )
            for i in order[:10]:
                lab = int(uniq_labels[i])
                r, g, b = (
                    [int(x) for x in label_rgb[lab]]
                    if 0 <= lab < len(label_rgb)
                    else (0, 0, 0)
                )
                print(f"  label={lab} rgb=({r},{g},{b}) faces={int(cnts[i])}")

        groups = defaultdict(list)
        for fi, lab in enumerate(face_label):
            groups[int(lab)].append(fi)

        for lab, face_ids in groups.items():
            if len(face_ids) < min_faces_per_part:
                continue

            sub = mesh.submesh([np.array(face_ids, dtype=np.int64)], append=True, repair=False)
            if sub is None:
                continue
            if isinstance(sub, (list, tuple)):
                if not sub:
                    continue
                sub = sub[0]

            if 0 <= lab < len(label_rgb):
                r, g, b = [int(x) for x in label_rgb[lab]]
                part_name = f"{base}__{node_name}__label_{lab}__rgb_{r}_{g}_{b}"
            else:
                part_name = f"{base}__{node_name}__label_{lab}"

            out_scene.add_geometry(sub, geom_name=part_name)
            part_count += 1

    if part_count == 0:
        if debug_print:
            print("[INFO] part_count==0, fallback to original scene export.")
        out_scene = scene

    out_scene.export(out_glb_path)
    return out_glb_path

def main():
    out_path = split_glb_by_texture_palette_rgb(
        INPUT_GLB,
        out_glb_path=None,
        min_faces_per_part=MIN_FACES_PER_PART,
        bake_transforms=BAKE_TRANSFORMS,
    )
    print("Done. Exported:", out_path)


if __name__ == "__main__":
    main()