import requests
import torch
import os

import trimesh.exchange
import trimesh.exchange.gltf
import folder_paths as comfy_paths
import logging
import time
import urllib.request
import datetime
import shutil
import glob
import json
import random
import time
import sys
from tqdm import tqdm
from server import PromptServer
from aiohttp import web
from PIL import Image, ImageOps, ImageSequence, ImageFile
from PIL.PngImagePlugin import PngInfo
import node_helpers
import numpy as np
import io
import comfy.utils
import trimesh

ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-Rodin")
CONFIG_ROOT_PATH = os.path.join(ROOT_PATH, "Configs")


# Constants
BASE_URL = "https://hyperhuman.deemos.com/api/v2"
API_HEADERS = {}

API_KEY_PARAMETER = {
    "api_key": ("APIKEY", {"forceInput": True, "multiline": True}),
}

COMMON_PARAMETERS = {
    "tier": (["Regular", "Sketch"], {"default": "Regular"}),
    "seed_": ("INT", {"default": 0, "min": 0, "max": 65535, "step": 1, "display": "number", }),
    "geometry_file_format":
    (["glb", "usdz", "fbx", "obj", "stl"],
     {"default": "glb", "label_off": "glb",
      "pysssss.binding": [{
          "source": "tier",
          "callback": [{
              "type": "if",
              "condition": [{"left": "$source.value", "op": "eq", "right": '"Sketch"'}],
              "true": [{
                  "type": "set", "target": "$this.options.values", "value": ["glb"]
              }, {
                  "type": "set", "target": "$this.value", "value": '"glb"'
              }, {
                  "type": "set", "target": "$this.disabled", "value": True
              },],
              "false": [{
                  "type": "set", "target": "$this.disabled", "value": False
              }, {
                  "type": "set", "target": "$this.options.values", "value": ["glb", "usdz", "fbx", "obj", "stl"]
              }],
          }]
      }]
      }),
    "material":
    (["PBR", "Shaded"],
     {"default": "PBR",
      "pysssss.binding": [{
          "source": "tier",
          "callback": [{
              "type": "if",
              "condition": [{"left": "$source.value", "op": "eq", "right": '"Sketch"'}],
              "true": [{
                  "type": "set", "target": "$this.options.values", "value": ["PBR"]
              }, {
                  "type": "set", "target": "$this.value", "value": '"PBR"'
              },
                  {
                  "type": "set", "target": "$this.disabled", "value": True
              },],
              "false": [{
                  "type": "set", "target": "$this.disabled", "value": False
              }, {
                  "type": "set", "target": "$this.options.values", "value": ["PBR", "Shaded"]
              }],
          }]
      }]
      }),
    "quality":
    (["high", "medium", "low", "extra-low"],
     {"default": "medium",
      "pysssss.binding": [{
          "source": "tier",
          "callback": [{
              "type": "if",
              "condition": [{"left": "$source.value", "op": "eq", "right": '"Sketch"'}],
              "true": [{
                  "type": "set", "target": "$this.options.values", "value": ["medium"]
              }, {
                  "type": "set", "target": "$this.value", "value": '"medium"'
              }, {
                  "type": "set", "target": "$this.disabled", "value": True
              },],
              "false": [{
                  "type": "set", "target": "$this.disabled", "value": False
              }, {
                  "type": "set", "target": "$this.options.values", "value": ["high", "medium", "low", "extra-low"]
              }],
          }]
      }]
      }),
    "use_hyper":
    ("BOOLEAN",
     {"default": False,
      "pysssss.binding": [{
          "source": "tier",
          "callback": [{
              "type": "if",
              "condition": [{"left": "$source.value", "op": "eq", "right": '"Sketch"'
                             }],
              "true": [{
                  "type": "set", "target": "$this.value", "value": False
              }, {
                  "type": "set", "target": "$this.disabled", "value": True
              },],
              "false": [{
                  "type": "set", "target": "$this.disabled", "value": False
              }, {
                  "type": "set", "target": "$this.options.values", "value": "$result"
              }],
          }]
      }]
      }),
}

SUPPORTED_3D_EXTENSIONS = (
    '.obj',
    '.glb',
    '.fbx',
    '.stl',
    '.usdz',
)


def post_request(url, api_key, data, files=None, max_retries=5, delay=2):
    headers = {"Authorization": f"Bearer {api_key}"}
    full_url = f"{BASE_URL}/{url}"

    for attempt in range(max_retries):
        try:
            response = requests.post(full_url, headers={**API_HEADERS, **headers}, data=data, files=files)
            response.raise_for_status()
            return response.json()
        except requests.ConnectionError as e:
            logging.info(f"Connection error: {e}. Retrying {attempt + 1}/{max_retries}...")
        except requests.Timeout as e:
            logging.info(f"Timeout error: {e}. Retrying {attempt + 1}/{max_retries}...")
        except requests.HTTPError as e:
            logging.info(f"HTTP error: {e}. Response: {response.text if response else 'No response'}")

            break
        except Exception as e:
            logging.info(f"An unexpected error occurred: {e}. Retrying {attempt + 1}/{max_retries}...")

        time.sleep(delay)
    
    logging.info("[ Rodin3D.process_request ] Max retries reached. Request failed.")
    return None
    


def check_status(api_key, subscription_key):
    data = {"subscription_key": subscription_key}
    return post_request("status", api_key, data)["jobs"]

def load_image(image_path):        
        img = node_helpers.pillow(Image.open, image_path)
        
        output_images = []
        w, h = None, None

        excluded_formats = ['MPO']
        
        for i in ImageSequence.Iterator(img):
            i = node_helpers.pillow(ImageOps.exif_transpose, i)

            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")

            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]
            
            if image.size[0] != w or image.size[1] != h:
                continue
            
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
            output_images.append(image)

        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
        else:
            output_image = output_images[0]

        return output_image

def handle_image(img):
    for i in ImageSequence.Iterator(img):
        i = node_helpers.pillow(ImageOps.exif_transpose, i)

        if i.mode == 'I':
            i = i.point(lambda i: i * (1 / 255))
        image = i.convert("RGB")
        w = image.size[0]
        h = image.size[1]
        
        if image.size[0] != w or image.size[1] != h:
            continue
        
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]
        if 'A' in i.getbands():
            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
            mask = 1. - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
        return image

def download_files(api_key, uuid):
    data = {"task_uuid": uuid}
    files_info = post_request("download", api_key, data)
    save_path = os.path.join(comfy_paths.get_output_directory(), datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))
    os.makedirs(save_path, exist_ok=True)
    
    shaded = diffuse = normal = pbr = load_image('custom_nodes/ComfyUI-Rodin/asset/error.png')
    for file_info in files_info["list"]:
        filename = file_info["name"].split('/')[-1]
        file_path = os.path.join(save_path, filename)
        if file_path.endswith(SUPPORTED_3D_EXTENSIONS):
            model_file_path = file_path
        print(f"[ download_files ] Downloading file: {file_path}")
        with requests.get(file_info["url"], stream=True) as r:
            if file_path.endswith('.glb'):
                kwargs = trimesh.exchange.gltf.load_glb(r.raw);
                shaded = handle_image(kwargs['geometry']['model']['visual'].material.baseColorTexture)
                diffuse = handle_image(kwargs['geometry']['model']['visual'].material.baseColorTexture)
                normal = handle_image(kwargs['geometry']['model']['visual'].material.normalTexture)
                pbr = handle_image(kwargs['geometry']['model']['visual'].material.metallicRoughnessTexture)
            with open(file_path, "wb") as f:
                shutil.copyfileobj(r.raw, f)
        if filename == "shaded.png":
            shaded = load_image(file_path)
        elif filename == "texture_diffuse.png":
            diffuse = load_image(file_path)
        elif filename == "texture_normal.png":
            normal = load_image(file_path)
        elif filename == "texture_pbr.png":
            pbr = load_image(file_path)


    #logging.info(model_file_path)

    return shaded, diffuse, normal, pbr, model_file_path, 


def tensor_to_filelike(tensor):
    """
    Converts a PyTorch tensor to a file-like object.

    Args:
    - tensor (torch.Tensor): A tensor representing an image of shape (H, W, C)
      where C is the number of channels (3 for RGB), H is height, and W is width.

    Returns:
    - io.BytesIO: A file-like object containing the image data.
    """
    array = tensor.cpu().numpy()
    array = (array * 255).astype('uint8')
    image = Image.fromarray(array, 'RGB')
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')  # PNG is used for lossless compression
    img_byte_arr.seek(0)
    return img_byte_arr

def LogInfomation(data, name):
    logging.info(f"[ Rodin3D.process_request ]\n{name} =")
    if isinstance(data, (dict)):
        logging.info(json.dumps(data, indent=4))
    elif isinstance(data, (list)) and name == "files":
        logging.info(repr(data))
    elif isinstance(data, (list)) and name == "status":
        table_header = f"{'UUID':<40} {'Status':<12}"
        separator = '-' * 52
        logging.info(table_header)
        logging.info(separator)
        for item in data:
            logging.info(f"{item['uuid']:<40} {item['status']:<12}")
        
    
    

class Rodin3D:

    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "STRING", )
    RETURN_NAMES = ("shaded", "diffuse", "normal", "pbr", "model_path")
    FUNCTION = "main_func"
    OUTPUT_NODE = True
    CATEGORY = "Mesh/Rodin"

    def process_request(self, api_key, images, prompt, condition_mode, seed, geometry_file_format, material, quality, use_hyper, tier) -> None:
        # Prepare request data and files
        files = [
            (
                "images",
                open(image, "rb") if isinstance(image, str) else tensor_to_filelike(image[0])
            )
            for image in images if image is not None]
        data = {
            "prompt": prompt,
            "condition_mode": condition_mode,
            "seed": seed,
            "geometry_file_format": geometry_file_format,
            "material": material,
            "quality": quality,
            "use_hyper": use_hyper,
            "tier": tier,
        }
        #logging.info(f"[ Rodin3D.process_request ]\n data = {data}, files = {files}")
        LogInfomation(data, "data")
        LogInfomation(files, "files")
        response = post_request("rodin", api_key, data, files=files)

        # Submit and handle the response
        if "uuid" in response:
            shaded, diffuse, normal, pbr, model_path = self.submit_poll_download(api_key, data, response['uuid'], response['jobs']['subscription_key'])
            return shaded, diffuse, normal, pbr, model_path, 
        else:
            logging.info(f"[ Rodin3D.process_request ] Error submitting the job:\n{response}")

    def submit_poll_download(self, api_key, data, uuid, subscription_key):
        """Submits the job, polls for its completion, and downloads the result when ready."""
        polling_interval = 2  # Interval in seconds to wait between checks

        total_seconds_estimated = 20 if data["tier"] == "Sketch" else 60
        pbar = comfy.utils.ProgressBar(total_seconds_estimated)

        while True:
            status = check_status(api_key, subscription_key)
            #logging.info(f"[ Rodin3D.process_request ] status = {status}")
            LogInfomation(status, "status")

            if all([job["status"] == "Done" for job in status]):
                logging.info(f"[ Rodin3D.process_request ] Generation complete. Downloading files...")
                save_model_path = download_files(api_key, uuid)
                print(f"[save model path] : {save_model_path}")
                return save_model_path

            time.sleep(polling_interval)
            pbar.update(2)


class RodinImage3D(Rodin3D):

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                **API_KEY_PARAMETER,
                "image": ("IMAGE", {"forceInput": True, "multiline": True}),
            },
            "optional": {
                "prompt": ("STRING",{"forceInput":True,"multiline": True}),
                **COMMON_PARAMETERS
            },
        }

    def main_func(self, api_key, image, seed_, geometry_file_format, material, quality, use_hyper, tier, prompt=None):
        images = [image]
        condition_mode = "concat"

        shaded, diffuse, normal, pbr, model_path = self.process_request(api_key, images, prompt, condition_mode, seed_, geometry_file_format, material, quality, use_hyper, tier)
        #logging.info(m_model_path)
        #logging.info(type(m_model_path))
        return (shaded, diffuse, normal, pbr, model_path)


class RodinMultipleImage3D(RodinImage3D):

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                **API_KEY_PARAMETER,
            },
            "optional": {
                "image_1": ("IMAGE", {"forceInput": True, "multiline": True}),
                "image_2": ("IMAGE", {"forceInput": True, "multiline": True}),
                "image_3": ("IMAGE", {"forceInput": True, "multiline": True}),
                "image_4": ("IMAGE", {"forceInput": True, "multiline": True}),
                "image_5": ("IMAGE", {"forceInput": True, "multiline": True}),
                "prompt": ("STRING",{"forceInput":True,"multiline": True}),
                "condition_mode": (["concat", "fuse"], {"default": "concat"}),
                **COMMON_PARAMETERS
            },
        }

    def main_func(self, api_key, seed_, geometry_file_format, material, quality, use_hyper, tier, condition_mode, image_1=None, image_2=None, image_3=None, image_4=None, image_5=None, prompt=None):
        images = [image_1, image_2, image_3, image_4, image_5]
        # Filter out None values
        images = [img for img in images if img is not None]

        model_path = self.process_request(api_key, images, prompt, condition_mode, seed_, geometry_file_format, material, quality, use_hyper, tier)

        return (model_path,)


class RodinText3D(RodinImage3D):

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                **API_KEY_PARAMETER,
                "prompt": ("STRING", {"forceInput": True, "multiline": True}),
            },
            "optional": {
                **COMMON_PARAMETERS
            },
        }

    def main_func(self, api_key, prompt, seed_, geometry_file_format, material, quality, use_hyper, tier):
        images = []
        condition_mode = None

        model_path = self.process_request(api_key, images, prompt, condition_mode, seed_, geometry_file_format, material, quality, use_hyper, tier)

        return (model_path,)


class PromptForRodin:

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("prompt",)

    FUNCTION = "main_func"
    CATEGORY = "Mesh/Rodin"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "prompt": ("STRING", {"default": "A textual prompt to guide the 3d generation.", "multiline": True})
            }
        }

    def main_func(self, prompt):
        return (prompt,)


class LoadRodinAPIKEY:

    RETURN_TYPES = ("APIKEY",)
    RETURN_NAMES = ("api_key",)
    FUNCTION = "main_func"
    CATEGORY = "Mesh/Rodin"

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "api_key": ("STRING", {"default": "Get your API KEY from: https://hyperhuman.deemos.com/api-dashboard", "multiline": True})
            },
        }

    def main_func(self, api_key):
        return (api_key,)



class Preview_3DMesh:

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_path":("STRING", {"default": '', "multiline": False,"forceInput": True, }),
                # NOTE: function which draging model file and getting the file path has done
                # "test": ("MODELUPLOAD", {"widget": "model_path"}),
            },
            # "optional": {
            #     "shaded":("IMAGE", {"multiline": True}),
            #     "diffuse":("IMAGE", { "multiline": True}),
            #     "normal":("IMAGE", {"multiline": True}),
            #     "pbr":("IMAGE", {"multiline": True}),
            # },
        }
    
    OUTPUT_NODE = True
    RETURN_TYPES = ()
    FUNCTION = "preview_mesh"
    CATEGORY = "Comfy3D/Visualize"

    def save_image(self, images, path):
        for (batch_number, image) in enumerate(images):
            i = 255. * image.cpu().numpy()
            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
            metadata = None
            img.save(path, pnginfo=metadata)
        return path

    
    def preview_mesh(self, model_path, shaded = None, diffuse = None, normal = None, pbr = None):
        
        model_folder, filename = os.path.split(model_path)

        
        if not os.path.isabs(model_path):
            model_folder = os.path.join(comfy_paths.output_directory, model_folder)
            model_path = os.path.join(comfy_paths.output_directory, model_path)
        
        if not filename.lower().endswith(SUPPORTED_3D_EXTENSIONS):
            logging.error(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3D file extensions: {SUPPORTED_3D_EXTENSIONS}")
            model_path = ""
        # shaded_path = diffuse_path = normal_path = pbr_path = ''

        # if shaded is not None:
        #     self.save_image(shaded, model_folder + '/shaded.png');
        #     shaded_path = model_folder + '/texture_shaded.png'
        # if diffuse is not None:
        #     self.save_image(diffuse, model_folder + '/diffuse.png');
        #     diffuse_path = model_folder + '/texture_diffuse.png'
        # if normal is not None:
        #     self.save_image(normal, model_folder + '/normal.png');
        #     normal_path = model_folder + '/texture_normal.png'
        # if pbr is not None:
        #     self.save_image(pbr, model_folder + '/pbr.png');
        #     pbr_path = model_folder + '/texture_pbr.png'
        
        previews = [
            {
                "model_path": model_path, 
                # "shaded_path": shaded_path, 
                # "diffuse_path": diffuse_path,
                # "normal_path": normal_path, 
                # "pbr_path": pbr_path, 
            }
        ]
        return {"ui": {"previews": previews}, "result": ()}


# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
    "RodinImage3D": RodinImage3D,
    "RodinMultipleImage3D": RodinMultipleImage3D,
    "RodinText3D": RodinText3D,
    "PromptForRodin": PromptForRodin,
    "LoadRodinAPIKEY": LoadRodinAPIKEY,
    "Preview_3DMesh":Preview_3DMesh,
}

# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
    "RodinImage3D": "Rodin - Image to 3D",
    "RodinMultipleImage3D": "Rodin - Multiple Images to 3D",
    "RodinText3D": "Rodin - Text to 3D",
    "PromptForRodin": "Rodin - Prompt for Rodin",
    "LoadRodinAPIKEY": "Rodin - API KEY",
    "Preview_3DMesh":"Rodin - Preview 3D Mesh",
}