"""Decima (simplifica) una malla 3D GLB/OBJ/PLY conservando su apariencia.

Post-proceso de las mallas densas que produce el pipeline Hunyuan3D de ComfyUI
(60 MB / ~1.67 M caras). Suelda los vertices duplicados ANTES del collapse —
las mallas de VoxelToMeshBasic son "cube-soup" (4 vertices nuevos por cara, sin
compartir aristas) y el quadric edge collapse no las decima si no se sueldan
primero. Luego aplica meshing_decimation_quadric_edge_collapse (pymeshlab), el
mismo filtro que el FaceReducer del wrapper Hy3D usa internamente.

Conserva la apariencia segun el tipo del original:
- Vertex colors (ColorVisuals, lo habitual en el shape stage de Hunyuan3D):
  pymeshlab los interpola a traves del collapse y se vuelven a adjuntar.
- Textura (TextureVisuals con UV + baseColorTexture): reproyecta las UV por
  vertice mas cercano (scipy cKDTree, sin rtree) y readjunta la imagen.
- Sin apariencia: simplifica la geometria igual.

Impura: lee y escribe archivos en disco. Requiere trimesh + pymeshlab + scipy.
"""
import os

import numpy as np


def _err(in_faces, in_mb, msg):
    return {
        "ok": False, "in_faces": in_faces, "out_faces": 0,
        "in_mb": in_mb, "out_mb": 0.0, "out_path": "",
        "welded_faces": 0, "decimate_status": "", "appearance": "",
        "error": msg,
    }


def _load_mesh(path):
    import trimesh

    obj = trimesh.load(path, process=False)
    if isinstance(obj, trimesh.Scene):
        obj = trimesh.util.concatenate(list(obj.geometry.values()))
    return obj


def _detect_appearance(mesh):
    """Devuelve (kind, attr, image): 'vertex_color'|'texture'|'none'."""
    import trimesh

    vis = mesh.visual
    if isinstance(vis, trimesh.visual.texture.TextureVisuals) and vis.uv is not None:
        img = getattr(getattr(vis, "material", None), "baseColorTexture", None)
        if img is not None:
            return "texture", np.asarray(vis.uv, dtype=np.float64), img
    if isinstance(vis, trimesh.visual.color.ColorVisuals):
        vc = vis.vertex_colors
        if vc is not None and len(vc) == len(mesh.vertices):
            return "vertex_color", np.asarray(vc, dtype=np.uint8), None
    return "none", None, None


def comfyui_simplify_mesh(
    in_path: str,
    *,
    target_faces: int = 80000,
    weld: bool = True,
    out_path: str | None = None,
) -> dict:
    """Decima una malla GLB/OBJ/PLY a target_faces conservando su apariencia.

    Args:
        in_path: ruta de la malla de entrada (.glb/.obj/.ply/.gltf/.stl/.off).
        target_faces: numero de caras objetivo tras la decimacion. keyword-only.
        weld: si True (default), suelda vertices duplicados antes del collapse.
            Imprescindible para las mallas cube-soup de VoxelToMeshBasic; sin el,
            el collapse no reduce caras. keyword-only.
        out_path: ruta de salida. Si None, escribe "<in>_simplified.glb" junto al
            original. keyword-only.

    Returns:
        dict {ok, in_faces, out_faces, in_mb, out_mb, out_path, welded_faces,
        decimate_status, appearance, error}. Si falla, ok=False y error explica.
    """
    try:
        import pymeshlab as ml
        import trimesh
    except ImportError as exc:
        return _err(0, 0.0, f"falta dependencia: {exc}. Instala en el venv del "
                            f"registry: cd python && uv add trimesh pymeshlab scipy")

    if not os.path.exists(in_path):
        return _err(0, 0.0, f"no existe el archivo de entrada: {in_path!r}")
    if out_path is None:
        out_path = os.path.splitext(in_path)[0] + "_simplified.glb"

    in_mb = round(os.path.getsize(in_path) / 1e6, 3)
    try:
        orig = _load_mesh(in_path)
    except Exception as exc:
        return _err(0, in_mb, f"no se pudo cargar la malla {in_path!r}: {exc}")

    in_faces = int(len(orig.faces))
    kind, attr, image = _detect_appearance(orig)

    try:
        ms = ml.MeshSet()
        if kind == "vertex_color":
            ms.add_mesh(ml.Mesh(
                vertex_matrix=orig.vertices, face_matrix=orig.faces,
                v_color_matrix=attr.astype(np.float64) / 255.0,
            ))
        else:
            ms.add_mesh(ml.Mesh(vertex_matrix=orig.vertices, face_matrix=orig.faces))

        if weld:
            ms.apply_filter("meshing_remove_duplicate_vertices")
            ms.apply_filter("meshing_remove_unreferenced_vertices")
        welded_faces = int(ms.current_mesh().face_number())

        decimate_status = "ok"
        if target_faces < welded_faces:
            ms.apply_filter(
                "meshing_decimation_quadric_edge_collapse",
                targetfacenum=int(target_faces), qualitythr=0.3,
                preserveboundary=True, boundaryweight=3.0,
                preservenormal=True, preservetopology=True, autoclean=True,
            )
        else:
            decimate_status = f"skipped_target_ge_welded({welded_faces})"

        cm = ms.current_mesh()
        dec = trimesh.Trimesh(
            vertices=cm.vertex_matrix(), faces=cm.face_matrix(), process=False,
        )

        appearance = kind
        if kind == "vertex_color" and cm.has_vertex_color():
            cols = np.clip(cm.vertex_color_matrix() * 255.0, 0, 255).astype(np.uint8)
            dec.visual = trimesh.visual.color.ColorVisuals(mesh=dec, vertex_colors=cols)
        elif kind == "texture":
            from scipy.spatial import cKDTree

            _, idx = cKDTree(orig.vertices).query(dec.vertices)
            dec.visual = trimesh.visual.texture.TextureVisuals(uv=attr[idx], image=image)

        parent = os.path.dirname(out_path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        dec.export(out_path)
    except Exception as exc:
        return _err(in_faces, in_mb, f"fallo al decimar/exportar: {type(exc).__name__}: {exc}")

    return {
        "ok": True,
        "in_faces": in_faces,
        "out_faces": int(len(dec.faces)),
        "in_mb": in_mb,
        "out_mb": round(os.path.getsize(out_path) / 1e6, 3),
        "out_path": out_path,
        "welded_faces": welded_faces,
        "decimate_status": decimate_status,
        "appearance": appearance,
        "error": "",
    }


if __name__ == "__main__":
    import json
    import sys

    src = sys.argv[1] if len(sys.argv) > 1 else (
        os.path.expanduser("~/ComfyUI/output/character_3d_00001_.glb"))
    tgt = int(sys.argv[2]) if len(sys.argv) > 2 else 80000
    out = sys.argv[3] if len(sys.argv) > 3 else None
    print(json.dumps(comfyui_simplify_mesh(src, target_faces=tgt, out_path=out), indent=2))