import base64
import io
import os
import struct
import tempfile

import numpy as np
import torch
from PIL import Image
from server import BinaryEventTypes, PromptServer

class Hy3DtoTD:
    #https://github.com/kijai/ComfyUI-Hunyuan3DWrapper
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "broadcast": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "trimesh": ("TRIMESH",),
                "mesh": ("MESH",),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_mesh"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_mesh(self, broadcast, trimesh=None, mesh=None):
        try:
            if trimesh is not None and mesh is not None:
                raise ValueError("Please connect only one input: trimesh or mesh, not both.")

            if trimesh is not None:
                vertices, colors = self.process_trimesh(trimesh)
            elif mesh is not None:
                vertices, colors = self.process_comfy_mesh(mesh)
            else:
                raise ValueError("Please connect either a trimesh or mesh input.")

            return self._send_to_td(vertices, colors, broadcast)

        except Exception as e:
            raise ValueError(f"Error processing mesh: {str(e)}")


    def process_trimesh(self, mesh):
        try:
            import trimesh

            if isinstance(mesh, trimesh.Scene):
                meshes = []
                for m in mesh.geometry.values():
                    if isinstance(m, trimesh.Trimesh):
                        meshes.append(m)

                if meshes:
                    mesh = trimesh.util.concatenate(meshes)
                else:
                    raise ValueError("No valid mesh data found")

            if not isinstance(mesh, trimesh.Trimesh):
                raise ValueError("Input data is not a valid Trimesh object")

            mesh.fix_normals()

            vertices = np.asarray(mesh.vertices, dtype=np.float32)
            vertex_count = len(vertices)

            if hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None:
                colors = np.asarray(mesh.visual.vertex_colors)
            elif hasattr(mesh.visual, "uv") and mesh.visual.uv is not None:
                try:
                    mesh.visual = mesh.visual.to_color()
                    colors = np.asarray(mesh.visual.vertex_colors)
                except Exception:
                    colors = np.tile([200, 200, 200, 255], (vertex_count, 1))
            else:
                colors = np.tile([200, 200, 200, 255], (vertex_count, 1))

            colors = self._normalize_colors(colors, vertex_count)

            return vertices, colors

        except ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error processing trimesh: {str(e)}")

    def process_comfy_mesh(self, mesh):
        try:
            if isinstance(mesh, list):
                if len(mesh) == 0:
                    raise ValueError("Empty mesh list received")
                mesh = mesh[0]

            vertices = self._extract_comfy_mesh_vertices(mesh)
            vertex_count = len(vertices)

            colors = self._extract_comfy_mesh_colors(mesh, vertex_count)
            colors = self._normalize_colors(colors, vertex_count)

            return vertices, colors

        except Exception as e:
            raise ValueError(f"Error processing ComfyUI mesh: {str(e)}")

    def _extract_comfy_mesh_vertices(self, mesh):
        vertices = None

        if hasattr(mesh, "vertices") and mesh.vertices is not None:
            vertices = mesh.vertices
        elif isinstance(mesh, dict) and "vertices" in mesh:
            vertices = mesh["vertices"]
        elif isinstance(mesh, (tuple, list)) and len(mesh) > 0:
            vertices = mesh[0]

        if vertices is None:
            raise ValueError("Cannot find vertices in ComfyUI mesh")

        vertices = self._to_numpy(vertices)
        vertices = self._squeeze_batch(vertices)
        vertices = np.asarray(vertices, dtype=np.float32)

        if vertices.ndim != 2 or vertices.shape[1] < 3:
            raise ValueError(f"Invalid vertices shape: {vertices.shape}")

        return vertices[:, :3]

    def _extract_comfy_mesh_colors(self, mesh, vertex_count):
        colors = None

        if hasattr(mesh, "vertex_colors") and mesh.vertex_colors is not None:
            colors = mesh.vertex_colors
        elif hasattr(mesh, "colors") and mesh.colors is not None:
            colors = mesh.colors
        elif hasattr(mesh, "visual") and mesh.visual is not None:
            if hasattr(mesh.visual, "vertex_colors") and mesh.visual.vertex_colors is not None:
                colors = mesh.visual.vertex_colors
        elif isinstance(mesh, dict) and "vertex_colors" in mesh:
            colors = mesh["vertex_colors"]
        elif isinstance(mesh, dict) and "colors" in mesh:
            colors = mesh["colors"]

        if colors is None:
            colors = np.tile([200, 200, 200, 255], (vertex_count, 1))

        return colors

    def _to_numpy(self, value):
        if value is None:
            return None

        if hasattr(value, "detach"):
            value = value.detach()
        if hasattr(value, "cpu"):
            value = value.cpu()
        if hasattr(value, "numpy"):
            value = value.numpy()

        return np.asarray(value)

    def _squeeze_batch(self, value):
        value = np.asarray(value)

        if value.ndim == 3 and value.shape[0] == 1:
            value = value[0]

        return value

    def _normalize_colors(self, colors, vertex_count):
        colors = self._to_numpy(colors)

        if colors is None:
            colors = np.tile([200, 200, 200, 255], (vertex_count, 1))

        colors = self._squeeze_batch(colors)
        colors = np.asarray(colors)

        if colors.ndim == 1:
            if colors.shape[0] in (3, 4):
                colors = np.tile(colors, (vertex_count, 1))
            else:
                colors = np.column_stack([colors, colors, colors])

        if colors.ndim != 2:
            colors = np.tile([200, 200, 200, 255], (vertex_count, 1))

        if len(colors) != vertex_count:
            if len(colors) > vertex_count:
                colors = colors[:vertex_count]
            elif len(colors) > 0:
                last_color = colors[-1]
                padding = np.tile(last_color, (vertex_count - len(colors), 1))
                colors = np.vstack([colors, padding])
            else:
                colors = np.tile([200, 200, 200, 255], (vertex_count, 1))

        if colors.dtype != np.uint8:
            if colors.size > 0 and colors.max() <= 1.0:
                colors = (np.clip(colors, 0, 1) * 255).astype(np.uint8)
            else:
                colors = np.clip(colors, 0, 255).astype(np.uint8)

        if colors.shape[1] == 1:
            colors = np.repeat(colors, 3, axis=1)

        if colors.shape[1] == 2:
            colors = np.repeat(colors[:, :1], 3, axis=1)

        if colors.shape[1] == 3:
            alpha = np.full((vertex_count, 1), 255, dtype=np.uint8)
            colors = np.concatenate([colors, alpha], axis=1)

        if colors.shape[1] > 4:
            colors = colors[:, :4]

        if colors.shape[1] < 4:
            default_colors = np.tile([200, 200, 200, 255], (vertex_count, 1)).astype(np.uint8)
            default_colors[:, :colors.shape[1]] = colors
            colors = default_colors

        return colors[:, :4].astype(np.uint8)

    def _send_to_td(self, vertices, colors, broadcast):
        try:
            vertices = np.asarray(vertices, dtype=np.float32)
            vertex_count = len(vertices)
            colors = self._normalize_colors(colors, vertex_count)

            dtype_vertex = np.dtype([
                ("pos", "<f4", 3),
                ("color", "u1", 4)
            ])

            vertex_data = np.empty(vertex_count, dtype=dtype_vertex)
            vertex_data["pos"] = vertices[:, :3]
            vertex_data["color"] = colors[:, :4]

            header = (
                "ply\n"
                "format binary_little_endian 1.0\n"
                f"element vertex {vertex_count}\n"
                "property float x\n"
                "property float y\n"
                "property float z\n"
                "property uchar red\n"
                "property uchar green\n"
                "property uchar blue\n"
                "property uchar alpha\n"
                "end_header\n"
            ).encode("ascii")

            binary_data = header + vertex_data.tobytes()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1000, binary_data, sid=sid)

            return {
                "ui": {
                    "mesh": [{
                        "source": "websocket",
                        "content-type": "model/ply",
                        "type": "output",
                    }]
                }
            }

        except Exception as e:
            raise ValueError(f"Error sending point PLY to TouchDesigner: {str(e)}")

class Tripo3DtoTD:
    #https://github.com/VAST-AI-Research/ComfyUI-Tripo
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_file": ("STRING",),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_mesh"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_mesh(self, model_file, broadcast):
        try:
            processed_mesh = self.process_mesh(model_file)
            return self._send_to_td(processed_mesh, broadcast)
        except Exception as e:
            raise ValueError(f"Error processing mesh: {str(e)}")

    def process_mesh(self, model_file):
        try:
            import trimesh
            if not model_file.startswith(('http://', 'https://')) and not os.path.isabs(model_file):
                try:
                    import folder_paths
                    output_dir = folder_paths.get_output_directory()
                    local_path = os.path.join(output_dir, model_file)
                    if os.path.exists(local_path):
                        model_file = local_path
                except:
                    pass

            mesh = trimesh.load(model_file)

            if isinstance(mesh, trimesh.Scene):
                meshes = []
                for m in mesh.geometry.values():
                    if isinstance(m, trimesh.Trimesh):
                        self._process_vertex_colors(m)
                        meshes.append(m)

                if meshes:
                    mesh = trimesh.util.concatenate(meshes)
                else:
                    raise ValueError("No valid mesh data found")

            if not isinstance(mesh, trimesh.Trimesh):
                raise ValueError("Input data is not a valid Trimesh object")

            mesh.fix_normals()

            if not hasattr(mesh.visual, 'vertex_colors') or mesh.visual.vertex_colors is None:
                mesh.visual = trimesh.visual.ColorVisuals(mesh)
                mesh.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(mesh.vertices), 1))
            else:
                if mesh.visual.vertex_colors.max() <= 1.0:
                    mesh.visual.vertex_colors = (mesh.visual.vertex_colors * 255).astype(np.uint8)
                else:
                    mesh.visual.vertex_colors = mesh.visual.vertex_colors.astype(np.uint8)

            return mesh

        except ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error processing mesh: {str(e)}")

    def _process_vertex_colors(self, m):
        try:
            import trimesh

            if not hasattr(m.visual, 'vertex_colors') or m.visual.vertex_colors is None:
                if hasattr(m.visual, 'material'):
                    material = m.visual.material
                    if hasattr(material, 'baseColorTexture') and material.baseColorTexture is not None:
                        if hasattr(m.visual, 'uv') and m.visual.uv is not None:
                            self._extract_colors_from_texture(m)
                    elif hasattr(material, 'baseColorFactor'):
                        color = np.array(material.baseColorFactor)
                        if len(color) == 3:
                            color = np.append(color, 1.0)
                        color = np.clip(color, 0, 1)
                        color = (color * 255).astype(np.uint8)
                        m.visual = trimesh.visual.ColorVisuals(m)
                        m.visual.vertex_colors = np.tile(color, (len(m.vertices), 1))
                    else:
                        m.visual = trimesh.visual.ColorVisuals(m)
                        m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
                else:
                    m.visual = trimesh.visual.ColorVisuals(m)
                    m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
            else:
                if m.visual.vertex_colors.max() <= 1.0:
                    m.visual.vertex_colors = (m.visual.vertex_colors * 255).astype(np.uint8)
                else:
                    m.visual.vertex_colors = m.visual.vertex_colors.astype(np.uint8)
        except ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error processing vertex colors: {str(e)}")

    def _extract_colors_from_texture(self, m):
        try:
            import trimesh

            material = m.visual.material
            image = material.baseColorTexture
            if image is not None:
                image_data = np.array(image)
                if image_data.dtype != np.uint8:
                    image_data = (np.clip(image_data, 0, 1) * 255).astype(np.uint8)

                uvs = m.visual.uv
                uvs = np.clip(uvs, 0, 1)
                h, w = image_data.shape[0], image_data.shape[1]
                i = ((1 - uvs[:, 1]) * (h - 1)).astype(int)
                j = (uvs[:, 0] * (w - 1)).astype(int)
                i = np.clip(i, 0, h - 1)
                j = np.clip(j, 0, w - 1)
                colors = image_data[i, j, :4]
                m.visual = trimesh.visual.ColorVisuals(m)
                m.visual.vertex_colors = colors
            else:
                m.visual = trimesh.visual.ColorVisuals(m)
                m.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(m.vertices), 1))
        except ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error extracting colors from texture: {str(e)}")

    def _send_to_td(self, mesh, broadcast):
        try:
            if not hasattr(mesh.visual, 'vertex_colors') or mesh.visual.vertex_colors is None:
                colors = np.tile([200, 200, 200, 255], (len(mesh.vertices), 1)).astype(np.uint8)
            else:
                colors = mesh.visual.vertex_colors.astype(np.uint8)

            header = [
                "ply",
                "format binary_little_endian 1.0",
                f"element vertex {len(mesh.vertices)}",
                "property float x",
                "property float y",
                "property float z",
                "property uchar red",
                "property uchar green",
                "property uchar blue",
                "property uchar alpha",
                "end_header\n"
            ]
            header = '\n'.join(header).encode('ascii')

            vertices = mesh.vertices.astype('<f4')
            colors = colors[:, :4].astype(np.uint8)

            vertex_count = len(vertices)
            binary_data = bytearray()
            for i in range(vertex_count):
                binary_data.extend(struct.pack('<fff', vertices[i, 0], vertices[i, 1], vertices[i, 2]))
                binary_data.extend(colors[i, :4].tobytes())

            buffer = io.BytesIO()
            buffer.write(header)
            buffer.write(binary_data)
            binary_output = buffer.getvalue()
            buffer.close()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1000, binary_output, sid=sid)

            return {"ui": {
                "mesh": [{
                    "source": "websocket",
                    "content-type": "model/ply",
                    "type": "output",
                }]
            }}
        except ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error sending mesh: {str(e)}")


class ImagetoTD:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_images"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_images(self, images, broadcast):
        results = []
        for tensor in images:
            array = 255.0 * tensor.cpu().numpy()
            image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(
                BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
                ["PNG", image, None],
                sid,
            )
            results.append({
                "source": "websocket",
                "content-type": "image/png",
                "type": "output",
            })

        return {"ui": {"images": results}}


class ImagetoTD_JPEG:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "broadcast": ("BOOLEAN", {"default": False}),
                "quality": ("INT", {"default": 95, "min": 1, "max": 100}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_images"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_images(self, images, broadcast, quality):
        results = []
        for tensor in images:
            array = 255.0 * tensor.cpu().numpy()
            image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))

            bytesIO = io.BytesIO()
            type_num = 1
            header = struct.pack(">I", type_num)
            bytesIO.write(header)

            image.save(bytesIO, format="JPEG", quality=quality)
            jpeg_bytes = bytesIO.getvalue()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id

            server.send_sync(
                BinaryEventTypes.PREVIEW_IMAGE,
                jpeg_bytes,
                sid,
            )
            results.append({
                "source": "websocket",
                "content-type": "image/jpeg",
                "type": "output",
            })

        return {"ui": {"images": results}}

class Comfy3DPacktoTD:
    #https://github.com/MrForExample/ComfyUI-3D-Pack
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mesh": ("MESH",),
                "broadcast": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "send_to_td"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_to_td(self, mesh, broadcast):
        try:
            vertices = mesh.v.cpu().numpy().astype(np.float32)
            colors = None
            color_source = "none"

            if hasattr(mesh, 'albedo') and mesh.albedo is not None:
                albedo = mesh.albedo
                if hasattr(albedo, 'cpu'):
                    albedo = albedo.cpu()
                if hasattr(albedo, 'numpy'):
                    albedo = albedo.numpy()

                if len(albedo.shape) > 1:
                    if hasattr(mesh, 'vt') and mesh.vt is not None:
                        uvs = mesh.vt.cpu().numpy()
                        uvs = np.clip(uvs, 0, 1)
                        tex_h, tex_w = albedo.shape[:2]

                        mapped_u, mapped_v = uvs[:, 0], uvs[:, 1]

                        tex_x = mapped_u * (tex_w - 1)
                        tex_y = mapped_v * (tex_h - 1)

                        tex_x0 = np.floor(tex_x).astype(np.int32)
                        tex_y0 = np.floor(tex_y).astype(np.int32)
                        tex_x1 = np.minimum(tex_x0 + 1, tex_w - 1)
                        tex_y1 = np.minimum(tex_y0 + 1, tex_h - 1)

                        wx = tex_x - tex_x0
                        wy = tex_y - tex_y0

                        c00 = albedo[tex_y0, tex_x0]
                        c01 = albedo[tex_y0, tex_x1]
                        c10 = albedo[tex_y1, tex_x0]
                        c11 = albedo[tex_y1, tex_x1]

                        c0 = c00 * (1 - wx[:, np.newaxis]) + c01 * wx[:, np.newaxis]
                        c1 = c10 * (1 - wx[:, np.newaxis]) + c11 * wx[:, np.newaxis]
                        colors = c0 * (1 - wy[:, np.newaxis]) + c1 * wy[:, np.newaxis]

                        color_source = "albedo_with_uv_no_v_flip"
                    else:
                        if albedo.shape[0] == len(vertices):
                            colors = albedo
                            color_source = "albedo_direct"
                else:
                    colors = albedo
                    color_source = "albedo_scalar"

            if colors is None and hasattr(mesh, 'vc') and mesh.vc is not None:
                colors = mesh.vc.cpu().numpy()
                color_source = "vertex_colors"

            if colors is None:
                positions = vertices - vertices.min(axis=0)
                positions = positions / positions.max(axis=0) if positions.max() > 0 else positions
                colors = positions
                color_source = "generated"

            if colors is not None:
                if colors.max() <= 1.0:
                    colors = (colors * 255).astype(np.uint8)
                else:
                    colors = np.clip(colors, 0, 255).astype(np.uint8)

            if colors.shape[1] == 3:
                alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
                colors = np.concatenate([colors, alpha], axis=1)

            header = [
                "ply",
                "format binary_little_endian 1.0",
                f"element vertex {len(vertices)}",
                "property float x",
                "property float y",
                "property float z",
                "property uchar red",
                "property uchar green",
                "property uchar blue",
                "property uchar alpha",
                "end_header\n"
            ]
            header = '\n'.join(header).encode('ascii')

            binary_data = bytearray()
            for i in range(len(vertices)):
                binary_data.extend(struct.pack('<fff',
                                               float(vertices[i, 0]), float(vertices[i, 1]), float(vertices[i, 2])))
                binary_data.extend(bytes([
                    int(colors[i, 0]),
                    int(colors[i, 1]),
                    int(colors[i, 2]),
                    int(colors[i, 3])
                ]))

            buffer = io.BytesIO()
            buffer.write(header)
            buffer.write(binary_data)
            binary_output = buffer.getvalue()
            buffer.close()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1000, binary_output, sid=sid)

            return {"ui": {
                "mesh": [{
                    "source": "websocket",
                    "content-type": "model/ply",
                    "type": "output",
                }]
            }}

        except Exception as e:
            print(f"[Error] send_to_td: {str(e)}")
            import traceback
            traceback.print_exc()
            raise ValueError(f"Error sending mesh: {str(e)}")


class TripoSRtoTD:
    #https://github.com/flowtyone/ComfyUI-Flowty-TripoSR
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "mesh": ("MESH",),
                "broadcast": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "send_to_td"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_to_td(self, mesh, broadcast):
        try:
            if isinstance(mesh, list):
                if len(mesh) > 0:
                    mesh_obj = mesh[0]
                else:
                    raise ValueError("Empty mesh list received")
            else:
                mesh_obj = mesh

            if hasattr(mesh_obj, 'vertices'):
                vertices = mesh_obj.vertices
                if hasattr(vertices, 'cpu'):
                    vertices = vertices.cpu()
                if hasattr(vertices, 'numpy'):
                    vertices = vertices.numpy()
                vertices = vertices.astype(np.float32)
            else:
                raise ValueError("Cannot find vertex data in mesh object")

            rotated_vertices = vertices.copy()
            rotated_vertices[:, 0] = -vertices[:, 0]
            rotated_vertices[:, 1], rotated_vertices[:, 2] = vertices[:, 2], vertices[:, 1]

            final_vertices = rotated_vertices.copy()
            final_vertices[:, 0] = rotated_vertices[:, 2]
            final_vertices[:, 2] = -rotated_vertices[:, 0]

            vertices = final_vertices

            colors = None

            if hasattr(mesh_obj, 'visual') and mesh_obj.visual is not None:
                visual = mesh_obj.visual

                if hasattr(visual, 'vertex_colors') and visual.vertex_colors is not None:
                    colors = visual.vertex_colors
                    if hasattr(colors, 'cpu'):
                        colors = colors.cpu()
                    if hasattr(colors, 'numpy'):
                        colors = colors.numpy()

                elif hasattr(visual, 'face_colors') and visual.face_colors is not None:
                    face_colors = visual.face_colors
                    if hasattr(face_colors, 'cpu'):
                        face_colors = face_colors.cpu()
                    if hasattr(face_colors, 'numpy'):
                        face_colors = face_colors.numpy()

                    if hasattr(mesh_obj, 'faces') and mesh_obj.faces is not None:
                        faces = mesh_obj.faces
                        if hasattr(faces, 'cpu'):
                            faces = faces.cpu()
                        if hasattr(faces, 'numpy'):
                            faces = faces.numpy()

                        colors = np.zeros((len(vertices), 4), dtype=np.uint8)
                        vertex_color_count = np.zeros(len(vertices), dtype=np.int32)

                        for i, face in enumerate(faces):
                            for vertex_idx in face:
                                colors[vertex_idx] += face_colors[i]
                                vertex_color_count[vertex_idx] += 1

                        for i in range(len(vertices)):
                            if vertex_color_count[i] > 0:
                                colors[i] = colors[i] // vertex_color_count[i]

                elif hasattr(visual, 'material') and visual.material is not None:
                    material = visual.material

                    if hasattr(material, 'diffuse') and material.diffuse is not None:
                        diffuse = material.diffuse
                        colors = np.tile(diffuse, (len(vertices), 1))

            if colors is None and hasattr(mesh_obj, '_visual') and mesh_obj._visual is not None:
                visual = mesh_obj._visual

                if hasattr(visual, 'vertex_colors') and visual.vertex_colors is not None:
                    colors = visual.vertex_colors
                    if hasattr(colors, 'cpu'):
                        colors = colors.cpu()
                    if hasattr(colors, 'numpy'):
                        colors = colors.numpy()

            if colors is None and hasattr(mesh_obj, 'visual') and hasattr(mesh_obj.visual,
                                                                          'uv') and mesh_obj.visual.uv is not None:
                if hasattr(mesh_obj.visual, 'material') and hasattr(mesh_obj.visual.material, 'image'):
                    uvs = mesh_obj.visual.uv
                    image = mesh_obj.visual.material.image

                    h, w = image.shape[:2]
                    u = np.clip(uvs[:, 0], 0, 1) * (w - 1)
                    v = np.clip(uvs[:, 1], 0, 1) * (h - 1)

                    u_int = np.round(u).astype(np.int32)
                    v_int = np.round(v).astype(np.int32)

                    u_int = np.clip(u_int, 0, w - 1)
                    v_int = np.clip(v_int, 0, h - 1)

                    colors = image[v_int, u_int]

            if colors is None:
                np.random.seed(42)
                colors = np.random.randint(100, 200, size=(len(vertices), 3), dtype=np.uint8)
                alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
                colors = np.concatenate([colors, alpha], axis=1)

            if colors is not None:
                if len(colors.shape) == 1:
                    colors = np.column_stack([colors, colors, colors])

                if colors.max() <= 1.0 and colors.dtype != np.uint8:
                    colors = (colors * 255).astype(np.uint8)
                else:
                    colors = np.clip(colors, 0, 255).astype(np.uint8)

                if colors.shape[1] == 3:
                    alpha = np.full((len(vertices), 1), 255, dtype=np.uint8)
                    colors = np.concatenate([colors, alpha], axis=1)
                elif colors.shape[1] > 4:
                    colors = colors[:, :4]

                if len(colors) != len(vertices):
                    if len(colors) > len(vertices):
                        colors = colors[:len(vertices)]
                    else:
                        last_color = colors[-1]
                        padding = np.tile(last_color, (len(vertices) - len(colors), 1))
                        colors = np.vstack([colors, padding])

            header = [
                "ply",
                "format binary_little_endian 1.0",
                f"element vertex {len(vertices)}",
                "property float x",
                "property float y",
                "property float z",
                "property uchar red",
                "property uchar green",
                "property uchar blue",
                "property uchar alpha",
                "end_header\n"
            ]
            header = '\n'.join(header).encode('ascii')

            binary_data = bytearray()
            for i in range(len(vertices)):
                binary_data.extend(struct.pack('<fff',
                                               float(vertices[i, 0]), float(vertices[i, 1]), float(vertices[i, 2])))
                binary_data.extend(bytes([
                    int(colors[i, 0]),
                    int(colors[i, 1]),
                    int(colors[i, 2]),
                    int(colors[i, 3])
                ]))

            buffer = io.BytesIO()
            buffer.write(header)
            buffer.write(binary_data)
            binary_output = buffer.getvalue()
            buffer.close()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1000, binary_output, sid=sid)

            return {"ui": {
                "mesh": [{
                    "source": "websocket",
                    "content-type": "model/ply",
                    "type": "output",
                }]
            }}

        except Exception as e:
            import traceback
            traceback.print_exc()
            raise ValueError(f"Error sending mesh: {str(e)}")


class LoadTDImage:
    # This node originates from the comfyui-tooling-nodes and utilizes the "Load Image (Base64)" functionality.
    # https://github.com/Acly/comfyui-tooling-nodes
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"image": ("STRING", {"multiline": False})}}

    RETURN_TYPES = ("IMAGE", "MASK")
    CATEGORY = "TouchDesigner"
    FUNCTION = "load_image"

    def load_image(self, image):
        try:
            imgdata = base64.b64decode(image)
            img = Image.open(io.BytesIO(imgdata))

            width, height = img.size

            if "A" in img.getbands():
                mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
                mask = 1.0 - torch.from_numpy(mask)
            else:
                mask = torch.zeros((height, width), dtype=torch.float32, device="cpu")

            if len(mask.shape) == 2:
                mask = mask.unsqueeze(0)

            img = img.convert("RGB")
            img = np.array(img).astype(np.float32) / 255.0
            img = torch.from_numpy(img)[None,]

            print(f"Image shape: {img.shape}, Mask shape: {mask.shape}")

            return (img, mask)

        except Exception as e:
            print(f"Error in LoadTDImage: {str(e)}")
            default_img = torch.zeros((1, 3, 64, 64), dtype=torch.float32, device="cpu")
            default_mask = torch.zeros((1, 64, 64), dtype=torch.float32, device="cpu")
            return (default_img, default_mask)


class VideotoTD:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "frame_rate": ("FLOAT", {"default": 8, "min": 1, "step": 1}),
                "quality": ("INT", {"default": 75, "min": 15, "max": 100, "step": 1}),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "process_video"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def process_video(self, images, frame_rate, broadcast, quality):
        if images is None or (isinstance(images, torch.Tensor) and images.size(0) == 0):
            raise ValueError("No valid images provided")

        video_binary = self._encode_video_ffmpeg(images, frame_rate, quality)

        return self._send_to_td(video_binary, broadcast)

    def _encode_video_ffmpeg(self, images, frame_rate, quality):
        try:
            import imageio_ffmpeg

            with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_file:
                temp_path = temp_file.name

            if isinstance(images, torch.Tensor):
                if len(images.shape) != 4:
                    raise ValueError(f"Expected 4D tensor, got shape {images.shape}")

                images_np = images.cpu().numpy()
                if images_np.max() <= 1.0:
                    images_np = (images_np * 255).astype(np.uint8)
                else:
                    images_np = images_np.astype(np.uint8)
            else:
                raise TypeError("Images must be a torch.Tensor")

            height, width = images_np[0].shape[:2]

            output_params = [
                '-c:v', 'libx264',
                '-preset', 'medium',
                '-crf', str(51 - quality // 2),
                '-pix_fmt', 'yuv420p'
            ]

            writer = imageio_ffmpeg.write_frames(
                temp_path,
                (width, height),
                fps=frame_rate,
                output_params=output_params
            )

            writer.send(None)

            for img in images_np:
                if img.shape[2] == 4:  # RGBA
                    img = img[:, :, :3]

                if not img.flags['C_CONTIGUOUS']:
                    img = np.ascontiguousarray(img)

                writer.send(img)

            writer.close()

            with open(temp_path, 'rb') as f:
                video_binary = f.read()

            os.unlink(temp_path)

            return video_binary

        except Exception as e:
            raise RuntimeError(f"Error encoding video with ffmpeg: {str(e)}")

    def _send_to_td(self, video_binary, broadcast):
        try:
            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1001, video_binary, sid=sid)

            return {"ui": {
                "video": [{
                    "source": "websocket",
                    "content-type": "video/mp4",
                    "type": "output",
                }]
            }}
        except Exception as e:
            raise ValueError(f"Error sending video: {str(e)}")
            
class AudiotoTD:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "audio": ("AUDIO",),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "process_audio"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def process_audio(self, audio, broadcast):
        if audio is None or not isinstance(audio, dict):
            raise ValueError("No valid audio dict provided")

        if "waveform" not in audio or "sample_rate" not in audio:
            raise ValueError("Audio dict must contain 'waveform' and 'sample_rate'")

        waveform = audio["waveform"]
        sample_rate = int(audio["sample_rate"])

        if not isinstance(waveform, torch.Tensor):
            raise TypeError("audio['waveform'] must be a torch.Tensor")

        if waveform.ndim != 3:
            raise ValueError(
                f"Expected waveform as 3D tensor [B, C, T], got shape {tuple(waveform.shape)}"
            )

        if waveform.size(0) > 1:
            waveform = waveform[0:1, ...]

        waveform = waveform[0]

        with torch.no_grad():
            max_abs = waveform.abs().amax()
            if torch.isfinite(max_abs) and max_abs > 1.0:
                waveform = waveform / max_abs

        wav_np = waveform.detach().cpu().float().numpy()

        if wav_np.ndim != 2:
            raise ValueError(
                f"Waveform must be [C, T] after squeeze, got {wav_np.shape}"
            )

        wav_np = np.transpose(wav_np, (1, 0))

        audio_binary = self._encode_wav(wav_np, sample_rate)

        return self._send_to_td(audio_binary, broadcast)

    def _encode_wav(self, wav_np, sample_rate):
        import wave

        if wav_np.ndim == 1:
            wav_np = wav_np[:, None]

        channels = wav_np.shape[1]

        pcm16 = np.clip(wav_np, -1.0, 1.0)
        pcm16 = (pcm16 * 32767.0).astype("<i2")

        if not pcm16.flags["C_CONTIGUOUS"]:
            pcm16 = np.ascontiguousarray(pcm16)

        buf = io.BytesIO()

        with wave.open(buf, "wb") as wf:
            wf.setnchannels(channels)
            wf.setsampwidth(2)
            wf.setframerate(sample_rate)
            wf.writeframes(pcm16.tobytes())

        audio_binary = buf.getvalue()
        buf.close()

        return audio_binary

    def _send_to_td(self, audio_binary, broadcast):
        server = PromptServer.instance
        sid = None if broadcast else server.client_id

        server.send_sync(1002, audio_binary, sid=sid)

        return {
            "ui": {
                "audio": [{
                    "source": "websocket",
                    "content-type": "audio/wav",
                    "type": "output",
                }]
            }
        }
        
class GaussianSplattingtoTD:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {},
            "optional": {
                "splat": ("SPLAT", {
                    "tooltip": "Optional input source. Choose either this SPLAT input or ply_path. If both are provided."
                }),
                "ply_path": ("STRING", {
                    "forceInput": True,
                    "default": "",
                    "tooltip": "Optional input source. Choose either a PLY file path or the SPLAT input. If both are provided,."
                }),
                "broadcast": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Send to all connected ComfyUI frontend clients. If disabled, the data is sent only to the current client."
                }),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_gaussian_splatting"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_gaussian_splatting(self, splat=None, ply_path="", broadcast=False):
        try:
            has_splat = splat is not None
            has_ply = bool(ply_path and ply_path.strip())

            if has_splat and has_ply:
                print(
                    "[GaussianSplattingtoTD] Both SPLAT and ply_path were provided. "
                    "Using SPLAT input and ignoring ply_path."
                )

            if has_splat:
                gs_data = self.process_splat(splat)
            elif has_ply:
                gs_data = self.process_ply(ply_path)
            else:
                raise ValueError(
                    "No input source provided. Please connect a SPLAT input or provide a ply_path. "
                    "Choose one of them."
                )

            return self._send_to_td(gs_data, broadcast)

        except Exception as e:
            raise ValueError(f"Error processing Gaussian Splatting data: {str(e)}")

    def process_splat(self, splat):
        def t2n(x):
            if x is None:
                return None
            if isinstance(x, torch.Tensor):
                return x.detach().cpu().float().numpy()
            return np.asarray(x, dtype=np.float32)

        positions = t2n(getattr(splat, 'positions'))
        scales    = t2n(getattr(splat, 'scales'))
        rotations = t2n(getattr(splat, 'rotations'))
        opacities = t2n(getattr(splat, 'opacities'))
        sh        = t2n(getattr(splat, 'sh'))

        counts = getattr(splat, 'counts', None)
        if counts is not None:
            counts = t2n(counts)

        if positions is None:
            raise ValueError("SPLAT is missing required field: positions")
        if scales is None:
            raise ValueError("SPLAT is missing required field: scales")
        if rotations is None:
            raise ValueError("SPLAT is missing required field: rotations")
        if opacities is None:
            raise ValueError("SPLAT is missing required field: opacities")
        if sh is None:
            raise ValueError("SPLAT is missing required field: sh")

        # Validate batched dimensions
        if positions.ndim != 3:
            raise ValueError(f"Expected positions shape (B, N, 3), got {positions.shape}")
        if scales.ndim != 3:
            raise ValueError(f"Expected scales shape (B, N, 3), got {scales.shape}")
        if rotations.ndim != 3:
            raise ValueError(f"Expected rotations shape (B, N, 4), got {rotations.shape}")
        if opacities.ndim != 3:
            raise ValueError(f"Expected opacities shape (B, N, 1), got {opacities.shape}")
        if sh.ndim != 4:
            raise ValueError(f"Expected sh shape (B, N, K, 3), got {sh.shape}")

        positions = positions[0]      # (N, 3)
        scales    = scales[0]         # (N, 3)
        rotations = rotations[0]      # (N, 4)
        opacities = opacities[0]      # (N, 1)
        sh        = sh[0]             # (N, K, 3)

        # Use counts[0] to remove padded points if available
        if counts is not None:
            n_real = int(counts[0])
            positions = positions[:n_real]
            scales    = scales[:n_real]
            rotations = rotations[:n_real]
            opacities = opacities[:n_real]
            sh        = sh[:n_real]

        f_dc = sh[:, 0, :].astype(np.float32)

        # Convert activated SPLAT values back to standard 3DGS PLY parameter space:
        eps = 1e-6

        op = np.clip(opacities.reshape(-1), eps, 1.0 - eps).astype(np.float32)
        opacity_logit = np.log(op / (1.0 - op))  # inverse sigmoid

        sc = np.clip(scales, eps, None).astype(np.float32)
        scale_log = np.log(sc)                   # inverse exp

        gs_data = {
            'positions': positions.astype(np.float32),
            'colors':    f_dc,
            'opacity':   opacity_logit,
            'scales':    scale_log,
            'rotations': rotations.astype(np.float32),
        }

        return gs_data

    def process_ply(self, ply_path):
        try:
            import os
            import numpy as np
            from plyfile import PlyData

            if not ply_path.startswith(('http://', 'https://')) and not os.path.isabs(ply_path):
                try:
                    import folder_paths
                    output_dir = folder_paths.get_output_directory()
                    local_path = os.path.join(output_dir, ply_path)
                    if os.path.exists(local_path):
                        ply_path = local_path
                except Exception:
                    pass

            if not os.path.exists(ply_path):
                raise ValueError(f"PLY file not found: {ply_path}")

            plydata = PlyData.read(ply_path)
            vertex = plydata['vertex']

            required_props = [
                'x', 'y', 'z',
                'f_dc_0', 'f_dc_1', 'f_dc_2',
                'opacity',
                'scale_0', 'scale_1', 'scale_2',
                'rot_0', 'rot_1', 'rot_2', 'rot_3'
            ]

            missing_props = [
                prop for prop in required_props
                if prop not in vertex.data.dtype.names
            ]

            if missing_props:
                raise ValueError(f"Missing properties in PLY file: {missing_props}")

            return {
                'positions': np.stack([
                    vertex['x'],
                    vertex['y'],
                    vertex['z']
                ], axis=1).astype(np.float32),

                'colors': np.stack([
                    vertex['f_dc_0'],
                    vertex['f_dc_1'],
                    vertex['f_dc_2']
                ], axis=1).astype(np.float32),

                'opacity': vertex['opacity'].astype(np.float32),

                'scales': np.stack([
                    vertex['scale_0'],
                    vertex['scale_1'],
                    vertex['scale_2']
                ], axis=1).astype(np.float32),

                'rotations': np.stack([
                    vertex['rot_0'],
                    vertex['rot_1'],
                    vertex['rot_2'],
                    vertex['rot_3']
                ], axis=1).astype(np.float32),
            }

        except ImportError:
            raise ImportError("Please install the plyfile library: pip install plyfile")
        except Exception as e:
            raise ValueError(f"Error reading PLY file: {str(e)}")

    def _send_to_td(self, gs_data, broadcast):
        try:
            import numpy as np
            from server import PromptServer

            num_points = len(gs_data['positions'])

            header = (
                "ply\n"
                "format binary_little_endian 1.0\n"
                f"element vertex {num_points}\n"
                "property float x\n"
                "property float y\n"
                "property float z\n"
                "property float f_dc_0\n"
                "property float f_dc_1\n"
                "property float f_dc_2\n"
                "property float opacity\n"
                "property float scale_0\n"
                "property float scale_1\n"
                "property float scale_2\n"
                "property float rot_0\n"
                "property float rot_1\n"
                "property float rot_2\n"
                "property float rot_3\n"
                "end_header\n"
            ).encode('ascii')

            # Vectorized packing:
            # 14 float32 values per point:
            # xyz + f_dc_rgb + opacity + scale_xyz + rotation_wxyz
            interleaved = np.concatenate([
                gs_data['positions'],                 # 3
                gs_data['colors'],                    # 3
                gs_data['opacity'].reshape(-1, 1),    # 1
                gs_data['scales'],                    # 3
                gs_data['rotations'],                 # 4
            ], axis=1).astype('<f4')

            binary_output = header + interleaved.tobytes()

            server = PromptServer.instance
            sid = None if broadcast else server.client_id

            server.send_sync(1003, binary_output, sid=sid)

            return {
                "ui": {
                    "gaussian_splatting": [{
                        "source": "websocket",
                        "content-type": "model/ply",
                        "type": "output",
                        "point_count": num_points,
                    }]
                }
            }

        except Exception as e:
            raise ValueError(f"Error sending Gaussian Splatting data: {str(e)}")

class StringtoTD:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "broadcast": ("BOOLEAN", {"default": False}),
                "prefix": ("STRING", {"default": ""}),
            },
            "optional": {
                "text": ("STRING",),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_string"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_string(self, broadcast, prefix, text=None):
        try:
            if text is None:
                text = ""

            final_text = str(prefix) + str(text)
            text_bytes = final_text.encode("utf-8")

            server = PromptServer.instance
            sid = None if broadcast else server.client_id
            server.send_sync(1004, text_bytes, sid=sid)

            return {
                "ui": {
                    "text": [{
                        "source": "websocket",
                        "content-type": "text/plain; charset=utf-8",
                        "type": "output",
                    }]
                }
            }
        except Exception as e:
            raise ValueError(f"Error sending string: {str(e)}")

class ImagetoTD_Tagged:
    EVENT_TYPE = 1005
    MAGIC = b"ITD3"
    VERSION = 1

    FORMAT_MAP = {
        "png": 1,
        "jpeg": 2,
    }

    CONTENT_TYPE_MAP = {
        "png": "image/png",
        "jpeg": "image/jpeg",
    }

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "format": (["png", "jpeg"], {"default": "png"}),
                "layer_tag": ("STRING", {"default": "depth", "multiline": False}),
                "quality": ("INT", {"default": 95, "min": 1, "max": 100}),
            }
        }

    RETURN_TYPES = ()
    FUNCTION = "send_images"
    OUTPUT_NODE = True
    CATEGORY = "TouchDesigner"

    def send_images(self, images, format, layer_tag, quality):
        results = []

        if layer_tag is None:
            layer_tag = ""
        layer_tag = str(layer_tag)

        for tensor in images:
            image = self._tensor_to_pil_image(tensor)
            image_bytes, content_type = self._encode_image(image, format, quality)
            packet = self._build_packet(format, layer_tag, image_bytes)

            server = PromptServer.instance

            server.send_sync(
                self.EVENT_TYPE,
                packet,
                sid=None,
            )

            results.append({
                "source": "websocket",
                "content-type": content_type,
                "type": "output",
                "layer-tag": layer_tag,
            })

        return {"ui": {"images": results}}

    def _tensor_to_pil_image(self, tensor):
        array = 255.0 * tensor.cpu().numpy()
        array = np.clip(array, 0, 255).astype(np.uint8)

        if array.ndim == 2:
            return Image.fromarray(array, mode="L")

        if array.ndim == 3:
            if array.shape[2] == 1:
                return Image.fromarray(array[:, :, 0], mode="L")
            elif array.shape[2] == 3:
                return Image.fromarray(array, mode="RGB")
            elif array.shape[2] == 4:
                return Image.fromarray(array, mode="RGBA")

        raise ValueError(f"Unsupported image tensor shape: {array.shape}")

    def _encode_image(self, image, format, quality):
        bytes_io = io.BytesIO()

        if format == "png":
            image.save(bytes_io, format="PNG")
            return bytes_io.getvalue(), self.CONTENT_TYPE_MAP["png"]

        elif format == "jpeg":
            if image.mode != "RGB":
                image = image.convert("RGB")
            image.save(bytes_io, format="JPEG", quality=quality)
            return bytes_io.getvalue(), self.CONTENT_TYPE_MAP["jpeg"]

        else:
            raise ValueError(f"Unsupported format: {format}")

    def _build_packet(self, format, layer_tag, image_bytes):
        """
        4s  magic      = b"ITD3"
        B   version    = 1
        B   format     = 1(png) / 2(jpeg)
        H   tag_len
        I   image_len
        ... tag utf-8 bytes
        ... image bytes
        """
        if format not in self.FORMAT_MAP:
            raise ValueError(f"Unsupported format: {format}")

        tag_bytes = layer_tag.encode("utf-8")
        fmt_code = self.FORMAT_MAP[format]

        if len(tag_bytes) > 65535:
            raise ValueError("layer_tag is too long, max 65535 bytes")

        header = struct.pack(
            ">4sBBHI",
            self.MAGIC,
            self.VERSION,
            fmt_code,
            len(tag_bytes),
            len(image_bytes),
        )

        return header + tag_bytes + image_bytes