import struct
import numpy as np
import io
from server import PromptServer, BinaryEventTypes
from pathlib import Path
from PIL import Image
import base64
from io import BytesIO
import torch
import tempfile
import os


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

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

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

    def process_mesh(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()

            if hasattr(mesh.visual, 'vertex_colors') and mesh.visual.vertex_colors is not None:
                pass
            elif hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
                mesh.visual = mesh.visual.to_color()
            else:
                mesh.visual = trimesh.visual.ColorVisuals(mesh)
                mesh.visual.vertex_colors = np.tile([200, 200, 200, 255], (len(mesh.vertices), 1))

            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 _send_to_td(self, mesh, broadcast):
        try:
            import trimesh

            buffer = io.BytesIO()
            mesh.export(file_obj=buffer, file_type='ply')
            binary_data = buffer.getvalue()
            buffer.close()

            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 ImportError:
            raise ImportError("Please install the trimesh library: pip install trimesh")
        except Exception as e:
            raise ValueError(f"Error sending mesh 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):
        import numpy as np
        from PIL import Image
        import io
        import struct

        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(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(s):
        return {
            "required": {
                "audio": ("AUDIO",),
                "format": (["mp3", "wav"],),
                "bitrate_kbps": ("INT", {"default": 192, "min": 32, "max": 512, "step": 16}),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

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

    def process_audio(self, audio, format, bitrate_kbps, 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))

        if format == "wav":
            audio_binary, content_type = self._encode_wav(wav_np, sample_rate)
        elif format == "mp3": 
            audio_binary, content_type = self._encode_mp3(wav_np, sample_rate, bitrate_kbps)
        else:
            raise ValueError(f"Unsupported format: {format}")

        return self._send_to_td(audio_binary, content_type, broadcast)

    def _encode_wav(self, wav_np, sample_rate):
        import soundfile as sf
        buf = io.BytesIO()
        if wav_np.ndim == 1:
            wav_np = wav_np[:, None]
        sf.write(buf, wav_np, samplerate=sample_rate, subtype="PCM_16", format="WAV")
        audio_binary = buf.getvalue()
        buf.close()
        return audio_binary, "audio/wav"

    def _encode_mp3(self, wav_np, sample_rate, bitrate_kbps):  
        import imageio_ffmpeg, subprocess
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as temp_file:  
            temp_path = temp_file.name
        try:
            channels = wav_np.shape[1] if wav_np.ndim == 2 else 1
            if wav_np.ndim == 1:
                wav_np = wav_np[:, None]
                channels = 1
            pcm = wav_np.astype(np.float32, copy=False)
            if not pcm.flags["C_CONTIGUOUS"]:
                pcm = np.ascontiguousarray(pcm)

            cmd = [
                imageio_ffmpeg.get_ffmpeg_exe(),
                "-y",
                "-f", "f32le",
                "-ac", str(channels),
                "-ar", str(sample_rate),
                "-i", "pipe:0",
                "-c:a", "libmp3lame", 
                "-b:a", f"{int(bitrate_kbps)}k",
                temp_path
            ]
            proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            proc.stdin.write(pcm.tobytes())
            proc.stdin.close()
            rc = proc.wait()
            if rc != 0:
                stderr = proc.stderr.read().decode("utf-8", errors="ignore")
                raise RuntimeError(f"ffmpeg failed (code {rc}): {stderr.strip()}")
            with open(temp_path, "rb") as f:
                audio_binary = f.read()
        finally:
            try:
                os.unlink(temp_path)
            except Exception:
                pass
        return audio_binary, "audio/mpeg"

    def _send_to_td(self, audio_binary, content_type, 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": content_type,
                "type": "output",
            }]
        }}

class GaussianSplattingtoTD:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ply_path": ("STRING",),
                "broadcast": ("BOOLEAN", {"default": False}),
            }
        }

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

    def send_gaussian_splatting(self, ply_path, broadcast):
        try:
            gs_data = self.process_ply(ply_path)
            return self._send_to_td(gs_data, broadcast)
        except Exception as e:
            raise ValueError(f"Error processing Gaussian Splatting data: {str(e)}")

    def process_ply(self, ply_path):
        try:
            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:
                    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}")
            
            gs_data = {
                '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),
            }
            
            return gs_data
            
        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:
            num_points = len(gs_data['positions'])
            
            header = [
                "ply",
                "format binary_little_endian 1.0",
                f"element vertex {num_points}",
                "property float x",
                "property float y",
                "property float z",
                "property float f_dc_0",
                "property float f_dc_1",
                "property float f_dc_2",
                "property float opacity",
                "property float scale_0",
                "property float scale_1",
                "property float scale_2",
                "property float rot_0",
                "property float rot_1",
                "property float rot_2",
                "property float rot_3",
                "end_header\n"
            ]
            header = '\n'.join(header).encode('ascii')
            
            binary_data = bytearray()
            for i in range(num_points):
                binary_data.extend(struct.pack('<fff', 
                    gs_data['positions'][i, 0], 
                    gs_data['positions'][i, 1], 
                    gs_data['positions'][i, 2]))
                
                binary_data.extend(struct.pack('<fff', 
                    gs_data['colors'][i, 0], 
                    gs_data['colors'][i, 1], 
                    gs_data['colors'][i, 2]))
                
                binary_data.extend(struct.pack('<f', gs_data['opacity'][i]))
                
                binary_data.extend(struct.pack('<fff', 
                    gs_data['scales'][i, 0], 
                    gs_data['scales'][i, 1], 
                    gs_data['scales'][i, 2]))
                
                binary_data.extend(struct.pack('<ffff', 
                    gs_data['rotations'][i, 0], 
                    gs_data['rotations'][i, 1], 
                    gs_data['rotations'][i, 2], 
                    gs_data['rotations'][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(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)}")