import trimesh
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:
    @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):
        processed_mesh = self.process_mesh(trimesh)
        return self._send_to_td(processed_mesh, broadcast)

    def process_mesh(self, mesh):
        try:
            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 Exception as e:
            raise ValueError(f"Error processing mesh: {str(e)}")

    def _send_to_td(self, mesh, broadcast):
        try:
            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 Exception as e:
            raise ValueError(f"Error sending mesh: {str(e)}")


class Tripo3DtoTD:
    @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):
        processed_mesh = self.process_mesh(model_file)
        return self._send_to_td(processed_mesh, broadcast)

    def process_mesh(self, model_file):
        try:
            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 Exception as e:
            raise ValueError(f"Error processing mesh: {str(e)}")

    def _process_vertex_colors(self, m):
        try:
            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 Exception as e:
            raise ValueError(f"Error processing vertex colors: {str(e)}")

    def _extract_colors_from_texture(self, m):
        try:
            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 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 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 Comfy3DPacktoTD:
    @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 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 = "external_tooling"
    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)}")