import os

from comfy.comfy_types.node_typing import IO
from comfy_extras.nodes_hunyuan3d import SaveGLB
from nodes import SaveImage


class BlenderOutputDownload3D():
    """Node used by ComfyUI Blender add-on to capture a 3D output from a workflow."""
    CATEGORY = "blender"
    FUNCTION = "execute"
    OUTPUT_NODE = True
    RETURN_TYPES = ()

    @classmethod
    def INPUT_TYPES(s):
        INPUT_TYPES = {"required": {}}
        INPUT_TYPES["required"]["model_file"] = (IO.STRING, {"forceInput": True})
        return INPUT_TYPES

    def execute(self, model_file: str) -> dict:

        # Get results
        subfolder, filename = os.path.split(model_file)
        result = {"filename": filename, "subfolder": subfolder, "type": "output"}

        # Push results
        websocket_message = {}
        websocket_message["ui"] = {}
        websocket_message["ui"]["3d"] = [result]
        return websocket_message

class BlenderOutputSaveGlb(SaveGLB):
    """Node used by ComfyUI Blender add-on to GLB file output from a workflow."""
    CATEGORY = "blender"

    @classmethod
    def INPUT_TYPES(s):
        INPUT_TYPES = super().INPUT_TYPES()
        INPUT_TYPES["required"]["filename_prefix"] = (IO.STRING, {"default": "3d/blender"})
        return INPUT_TYPES

class BlenderOutputSaveImage(SaveImage):
    """Node used by ComfyUI Blender add-on to capture an image output from a workflow."""
    CATEGORY = "blender"

    @classmethod
    def INPUT_TYPES(s):
        INPUT_TYPES = super().INPUT_TYPES()
        INPUT_TYPES["required"]["filename_prefix"] = (IO.STRING, {"default": "blender"})
        return INPUT_TYPES