"""
Main node definitions for the TensorRT_RTX upscaling extension.

This module provides ComfyUI nodes for building, loading, and executing 
TensorRT-optimized upscaling engines. It includes support for dynamic 
shape configurations and automatic engine compilation from ONNX models.
"""
import os
import time
import gc
import logging

import numpy as np
import tensorrt_rtx as trt
import torch

import comfy.model_management as mm
import folder_paths

from .utilities import Engine
from comfy.utils import ProgressBar
from spandrel import ModelLoader, ImageModelDescriptor


IMAGE_DIM_MIN = 256
IMAGE_DIM_OPT = 512
IMAGE_DIM_MAX = 1280

IMAGE_BATCH_MIN = 1
IMAGE_BATCH_MAX = 1

def update_folder_names_and_paths(key, targets=[]):
    """
    Ensures specific folder paths are registered within ComfyUI's path manager.

    Args:
        key (str): The key to register in folder_names_and_paths.
        targets (list): Fallback keys to inherit paths from if the primary key is missing.
    """
    # check for existing key
    base = folder_paths.folder_names_and_paths.get(key, ([], {}))
    base = base[0] if isinstance(base[0], (list, set, tuple)) else []
    # find base key & add w/ fallback, sanity check + warning
    target = next((x for x in targets if x in folder_paths.folder_names_and_paths), targets[0])
    orig, _ = folder_paths.folder_names_and_paths.get(target, ([], {}))
    folder_paths.folder_names_and_paths[key] = (orig or base, {".onnx"})
    if base and base != orig:
        logging.warning(f"Unknown file list already present on key {key}: {base}")

# Add a custom keys for files ending in .onnx
update_folder_names_and_paths("onnx_upscalers", ["onnx"])

class TensorRTXUpscaleImage:
    """
    A ComfyUI node that performs image upscaling using a loaded TensorRT engine.
    """    
    @classmethod    
    def INPUT_TYPES(cls):
        return {
                "required": {
                    "images": ("IMAGE", {"tooltip": "Images to be upscaled. Resolution must be between 256 and 1280 px if no build options has been specified."}),
                    "engine": ("TRT_RTX_ENGINE", {"tooltip": "TensorRT_RTX engine built and loaded"}),
                    "scale_by": (["2x", "3x", "4x", "custom"], {"default": "2x", "tooltip": "Target upscaling factor"})                    
                },
                "optional": {
                    "width": ("INT", {"default": 1024, "min": 256, "max": 4096}),
                    "height": ("INT", {"default": 1024, "min": 256, "max": 4096}),
                },
            }
   
    RETURN_NAMES = ("upscaled_images",)
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "upscale"
    CATEGORY = "TensorRT_RTX/Upscale"
    DESCRIPTION = "Upscale images with TensorRT_RTX"

    def upscale(self, **kwargs):
        """
        Executes the upscaling inference on the provided images.

        Args:
            image (torch.Tensor): Input image tensor in NHWC format.
            engine (TRT_RTX_ENGINE): The Tensorrt_RTX object.
            resize (dict, optional): Desired output dimensions. Defaults to 4x scale.

        Returns:
            tuple[torch.Tensor]: A tuple containing the upscaled image tensor in NHWC format.
        """
        images      = kwargs.get("images") 
        engine      = kwargs.get("engine")
        scale_by    = kwargs.get("scale_by")
        images_bchw = images.permute(0, 3, 1, 2)
        B, C, H, W  = images_bchw.shape

        # Bound checking
        for dim in (H,W):
            if not IMAGE_DIM_MIN <= dim <= IMAGE_DIM_MAX:
                raise ValueError(f"The input image dimension ({dim}) is outside of Tensorrt_RTX engine bounds: {IMAGE_DIM_MIN} to {IMAGE_DIM_MAX}.\nImage Dimensions:({W}x{H})")

        final_width, final_height = None, None

        if scale_by == "custom":
            final_width  = kwargs.get("width")
            final_height = kwargs.get("height")
        else:
            resize_factor = float(scale_by.split('x')[0])
            final_width   = int(W * resize_factor)
            final_height  = int(H * resize_factor)
            
        logging.info(f"Upscaling {B} images from H:{H}, W:{W} to H:{H*4}, W:{W*4} | Final resolution: H:{final_height}, W:{final_width}")

        shape_dict = {
            "input": {"shape": (1, 3, H, W)},
            "output": {"shape": (1, 3, H*4, W*4)},
        }
        
        # setup engine
        engine.activate()
        engine.allocate_buffers(shape_dict=shape_dict)

        cudaStream      = torch.cuda.current_stream().cuda_stream
        pbar            = ProgressBar(B)
        images_list     = list(torch.split(images_bchw, split_size_or_sections=1))
        upscaled_frames = torch.empty((B, C, final_height, final_width), dtype=torch.float32, device=mm.intermediate_device()) # offloaded to cpu

        for i, img in enumerate(images_list):
            result = engine.infer({"input": img}, cudaStream)
            result = result["output"]
            
            if W*4 != final_width or H*4 != final_height:
                if W*4 > final_width and H*4 > final_height:
                    mode='area'
                    antialias=False
                else:
                    mode='bicubic'
                    antialias=True
                    
                result = torch.nn.functional.interpolate(result, size=(final_height, final_width), mode=mode, antialias=antialias)
                
            upscaled_frames[i] = result.to(mm.intermediate_device())
            pbar.update(1)

        output = upscaled_frames.permute(0, 2, 3, 1)
        
        #Free VRAM
        engine.reset()
        gc.collect()
        mm.soft_empty_cache()

        logging.info(f"Output shape: {output.shape}")
        return (output,)

class TensorRTXEngineLoader:
    """
    A ComfyUI node that loads a TensorRT engine, building it from ONNX if necessary.
    """
    @classmethod
    def INPUT_TYPES(cls):
        onnx_names = [x for x in folder_paths.get_filename_list("onnx_upscalers") if "rife" not in x.lower()]        
        return {
            "required": { "onnx_name": (onnx_names,)},
            "optional": {
                "config": ("TRT_RTX_ENGINE_CONFIG", {"tooltip": "Options for building the TensorRT engine"}),
            }
        }
        
    RETURN_NAMES = ("ENGINE",)
    RETURN_TYPES = ("TRT_RTX_ENGINE",)
    CATEGORY = "TensorRT_RTX/Upscaler"
    DESCRIPTION = "Load TensorRT model (the model will be built automatically if not found)"
    FUNCTION = "load_upscaler_tensorrt_model"

    def load_upscaler_tensorrt_model(self, onnx_name, config=None):
        """
        Loads the .trt engine or compiles it from .onnx using specified dynamic shapes.

        Args:
            onnx_name (str): The filename of the source ONNX model.
            config (TensorRTXDynamicShapeOptions, optional): Configuration for dynamic shapes.

        Returns:
            engine [TRT_ENGINE]: Contains the 'engine' instance.
        """
        if config is None:
            config = TensorRTXDynamicShapeOptions()

        tensorrt_models_dir = os.path.join(folder_paths.models_dir, "tensorrt", "upscaler")
        onnx_models_dir     = os.path.join(folder_paths.models_dir, "onnx")

        os.makedirs(tensorrt_models_dir, exist_ok=True)
        os.makedirs(onnx_models_dir, exist_ok=True)

        onnx_model_path = os.path.join(onnx_models_dir, f"{onnx_name}")

        # Dynamic shape configuration
        engine_channel = 3
        engine_min_batch, engine_opt_batch, engine_max_batch = config.batch_min, config.batch_opt, config.batch_max
        engine_min_h, engine_opt_h, engine_max_h = config.height_min, config.height_opt, config.height_max
        engine_min_w, engine_opt_w, engine_max_w = config.width_min, config.width_opt, config.width_max
        tensorrt_model_path = os.path.join(tensorrt_models_dir, f"{onnx_name}_{engine_min_batch}x{engine_channel}x{engine_min_h}x{engine_min_w}_{engine_opt_batch}x{engine_channel}x{engine_opt_h}x{engine_opt_w}_{engine_max_batch}x{engine_channel}x{engine_max_h}x{engine_max_w}_{trt.__version__}.trt")

        # Check if the TRT file exists and build it if it doesn't
        engine = Engine(tensorrt_model_path)       
        
        if not os.path.exists(tensorrt_model_path): 
            logging.info(f"Building TensorRT_RTX engine for {onnx_model_path}: {tensorrt_model_path}")
            
            mm.soft_empty_cache()
            s      = time.time()            
            result = engine.build(
                                  onnx_path=onnx_model_path,                                  
                                  input_profile=[
                                      {"input": [
                                          (engine_min_batch,engine_channel,engine_min_h,engine_min_w),
                                          (engine_opt_batch,engine_channel,engine_opt_h,engine_min_w), 
                                          (engine_max_batch,engine_channel,engine_max_h,engine_max_w)
                                      ]}, 
                                  ]
            )
            if result != 0:
                raise ValueError(f"Failed to build the engine")
            e = time.time()
            logging.info(f"Time taken to build: {(e-s)} seconds")

        # Load tensorrt model
        logging.info(f"Loading TensorRT_RTX engine: {tensorrt_model_path}")
        mm.soft_empty_cache()
        engine.load()

        return (engine,)

class TensorRTXDynamicShapeOptions:
    """
    Data container for TensorRT dynamic shape boundaries and optimal values.
    """
    width_min: int = 256
    width_opt: int = 512
    width_max: int = 1280
    height_min: int = 256
    height_opt: int = 512
    height_max: int = 1280
    batch_min: int = 1
    batch_opt: int = 1
    batch_max: int = 1


class TensorRTXDynamicShapeConfig:
    """
    A ComfyUI node that allows users to configure TensorRT dynamic shape profiles.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "width_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_min, "min": 1, "max": 4096}),
                "width_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_opt, "min": 1, "max": 4096}),
                "width_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.width_max, "min": 1, "max": 4096}),
                "height_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_min, "min": 1, "max": 4096}),
                "height_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_opt, "min": 1, "max": 4096}),
                "height_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.height_max, "min": 1, "max": 4096}),
                "batch_min" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_min, "min": 1}),
                "batch_opt" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_opt, "min": 1}),
                "batch_max" : ("INT", {"default": TensorRTXDynamicShapeOptions.batch_max, "min": 1}),
            },
        }

    RETURN_TYPES = ("TRT_RTX_ENGINE_CONFIG",)
    RETURN_NAMES = ("CONFIG",)
    CATEGORY = "TensorRT_RTX/Upscaler"
    FUNCTION = "package"

    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """
        Validates that min <= opt <= max for all shape dimensions.
        """
        if not all(v is not None for v in kwargs.values()):
            raise ValueError(f"Dynamic shape properties must be numbers")
        
        # Check dimensions logic (min <= opt <= max)        
        for dim in ['width', 'height', 'batch']:
            mi, opt, ma = kwargs[f'{dim}_min'], kwargs[f'{dim}_opt'], kwargs[f'{dim}_max']
            if not 1 <= mi <= opt <= ma:
                raise ValueError(f"The {dim} bounds are invalid. Must follow: 1 <= min <= opt <= max")
                
        return True

    def package(self, **kwargs):
        """
        Instantiates and populates a TensorRTXDynamicShapeOptions object.
        """
        
        options = TensorRTXDynamicShapeOptions()
        
        for key, value in kwargs.items():
            setattr(options, key, value)
        return (options,)

# Node Mappings
NODE_CLASS_MAPPINGS = {
    "TensorRTXUpscaleImage": TensorRTXUpscaleImage,
    "TensorRTXEngineLoader": TensorRTXEngineLoader,
    "TensorRTXDynamicShapeConfig": TensorRTXDynamicShapeConfig,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TensorRTXUpscaleImage": "RTX Upscale Image",
    "TensorRTXEngineLoader": "RTX Engine Loader",
    "TensorRTXDynamicShapeConfig": "RTX Engine Dynamic Shape Config",
}