import node_helpers
import os

from .helpers.dense_diffusion import dd_combine, dd_apply
from .helpers.ipadapter import apply_ipadapter
from .helpers.imgutil import create_solid_mask
from .helpers.imgutil import scale_tensor_image

class PseudoApplyDenseDiffusionSDXL:
    """
    Processor class for applying dense diffusion to material prompts within a scene context.
    Inputs:
        model (MODEL): The base model to apply dense diffusion to.
        clip (CLIP): The CLIP model used for text encoding.
        mat_txts_lst (list of str): List of material prompt texts, one for each object/material.
        mat_msks_lst (list of tensor): List of mask tensors as [1, H, W], one for each material prompt.
        env_scene (str): Scene description prompt.
        env_style (str): Style description prompt to be appended to each material prompt.
        env_negative (str): Negative prompt for conditioning.
        width (int): Target width for all masks and outputs.
        height (int): Target height for all masks and outputs.
    Outputs:
        model (MODEL): The processed model after applying dense diffusion with all prompts and masks.
        positive (CONDITIONING): The positive conditioning tensor for the combined scene and material prompts.
        negative (CONDITIONING): The negative conditioning tensor for the negative prompt.
    Additional Information:
        - Material masks are automatically resized to the specified width and height if needed.
        - Each material prompt is combined with the style prompt for conditioning.
    """
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL", {"forceInput": True}),
                "clip": ("CLIP", {"forceInput": True}),
                "mat_txts_lst": ("STRING", {"forceInput": True}),
                "mat_msks_lst": ("MASK", {"forceInput": True}),
                "env_scene": ("STRING", {"forceInput": True}),
                "env_style": ("STRING", {"forceInput": True}),
                "env_negative": ("STRING", {"forceInput": True}),
                "width": ("INT", {"forceInput": True}),
                "height": ("INT", {"forceInput": True}),
            },
        }
    
    INPUT_IS_LIST = True # All inputs of ``type`` will become ``list[type]``, regardless of how many items are passed in.

    RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING",)
    RETURN_NAMES = ("model", "positive", "negative",)
    OUTPUT_IS_LIST = (False, False, False,)

    FUNCTION = "func"

    CATEGORY = "Pseudocomfy/Conditioning"

    def func(self, model, clip, mat_txts_lst, mat_msks_lst, env_scene, env_style, env_negative, width, height):
        # material inputs are expected to be lists
        mat_txts = mat_txts_lst 
        mat_msks = mat_msks_lst

        # if model or clip is a list, use the first element
        if isinstance(model, list) and len(model)>0: model = model[0]
        if isinstance(clip, list) and len(clip)>0: clip = clip[0]

        # If env_x are a list, concatenate into a single string
        if isinstance(env_scene, list): env_scene = ", ".join(env_scene)
        if isinstance(env_style, list): env_style = ", ".join(env_style)
        if isinstance(env_negative, list): env_negative = ", ".join(env_negative)

        # if width or height are a list, use the first element
        if isinstance(width, list) and len(width)>0: width = width[0]
        if isinstance(height, list) and len(height)>0: height = height[0]
        
        # Create a report string pairing each mat_txt with its corresponding mat_msk shape
        mat_report = "\n".join(
            f"\tmat[{i}]: {tuple(msk.shape)} '{txt[:20]}...'"
            for i, (txt, msk) in enumerate(zip(mat_txts, mat_msks))
        )
        print(f"[pseudocomfy] ApplyDenseDiffusion\n\tenv_scene: '{env_scene[:20]}...'\n\tenv_style: '{env_style[:20]}...'\n\tenv_negative: '{env_negative[:20]}...'\n\twidth, height: ({width},{height})\n{mat_report}")

        # Ensure all masks have shape [1, width, height]
        for i in range(len(mat_msks)):
            mask = mat_msks[i]
            # Check mask has three dimensions and the first dimension is 1
            if not (isinstance(mask.shape, tuple) and len(mask.shape) == 3 and mask.shape[0] == 1):
                raise ValueError(f"Mask at index {i} must have shape [1, width, height], got {mask.shape}")            
            
            # scale the mask to the desired width and height if necessary
            if mask.shape[-2] != width or mask.shape[-1] != height:
                #print("[pseudocomfy] ApplyDenseDiffusion\n\tscaling mask to width, height: ", width, height)
                mat_msks[i] = scale_tensor_image(mask, width, height)

        styled_material_prompts = [prompt + ", " + env_style for prompt in mat_txts] # adding styles to each object prompt
        # turning the list of strings into a list of conditionings:
        mat_pmts_cond = [clip_text_encode(clip, prompt) for prompt in styled_material_prompts] # appending as a list - format of comfy when returning CONDITIONING type

        combined_txt_list = [env_scene, env_style] + mat_txts # list containing all scene, style and object prompts 
        env_positive = "; ".join(combined_txt_list) # combine all prompts into a single string

        positive_prompt_cond = clip_text_encode(clip, env_positive) # wrapping in a list - format of comfy when returning CONDITIONING type
        negative_prompt_cond = clip_text_encode(clip, env_negative)

        empty_mask = create_solid_mask(1.0, width, height)
        model = dd_combine(model, positive_prompt_cond, empty_mask, 1.0) # first combining with dense diffusion

        for i in range(len(mat_pmts_cond)):
            model = dd_combine(model, mat_pmts_cond[i], mat_msks[i], 1.0)

        work_model, cond = dd_apply(model)

        return (work_model, cond, negative_prompt_cond)
    

class PseudoApplyIPAdaperSDXL:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "ipadapter": ("IPADAPTER",),
                "mat_txts_lst": ("STRING", {"forceInput": True}), # we expect a list of strings
                "mat_imgs_lst": ("IMAGE", {"forceInput": True}), # we expect a list of images
                "mat_msks_lst": ("MASK", {"forceInput": True}), # we expect a list of masks
                "env_scene": ("STRING", {"forceInput": True}),
                "env_style": ("STRING", {"forceInput": True}),
                "env_negative": ("STRING", {"forceInput": True}),
                "width": ("INT", {"forceInput": True}),
                "height": ("INT", {"forceInput": True}),
                "env_cond_strength": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 10.0, "step": 0.01}),
                "mat_cond_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
                "ipadapter_weight": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.05 }),
                "start_at": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.05 }),
                "end_at": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05 }),
            },
        }
    
    INPUT_IS_LIST = True # All inputs of ``type`` will become ``list[type]``, regardless of how many items are passed in.

    RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING",)
    RETURN_NAMES = ("model", "positive", "negative",)
    OUTPUT_IS_LIST = (False, False, False,)

    FUNCTION = "func"
    CATEGORY = "Pseudocomfy/Conditioning"

    def func(self, model, clip, ipadapter, mat_txts_lst, mat_imgs_lst, mat_msks_lst, env_scene, env_style, env_negative, width, height, env_cond_strength, mat_cond_strength, ipadapter_weight, start_at, end_at):
        # material inputs are expected to be lists
        mat_txts = mat_txts_lst 
        mat_imgs = mat_imgs_lst
        mat_msks = mat_msks_lst

        # unwrap all expected singletons
        #

        
        # if model or clip or ipadapter is a list, use the first element
        if isinstance(model, list) and len(model)>0: model = model[0]
        if isinstance(clip, list) and len(clip)>0: clip = clip[0]
        if isinstance(ipadapter, list) and len(ipadapter)>0: ipadapter = ipadapter[0]

        # If env_x are a list, concatenate into a single string
        if isinstance(env_scene, list): env_scene = ", ".join(env_scene)
        if isinstance(env_style, list): env_style = ", ".join(env_style)
        if isinstance(env_negative, list): env_negative = ", ".join(env_negative)

        # if width or height are a list, use the first element
        if isinstance(width, list) and len(width)>0: width = width[0]
        if isinstance(height, list) and len(height)>0: height = height[0]     
        
        # if any strength is a list, use the first element
        if isinstance(env_cond_strength, list) and len(env_cond_strength)>0: env_cond_strength = env_cond_strength[0]
        if isinstance(mat_cond_strength, list) and len(mat_cond_strength)>0: mat_cond_strength = mat_cond_strength[0]
        if isinstance(ipadapter_weight, list) and len(ipadapter_weight)>0: ipadapter_weight = ipadapter_weight[0]

        # start_at and end_at are singletons
        if isinstance(start_at, list) and len(start_at)>0: start_at = start_at[0]
        if isinstance(end_at, list) and len(end_at)>0: end_at = end_at[0]

        # report inputs
        #

        print(f"[pseudocomfy] ApplyIPAdaper\n\tenv_scene: '{env_scene[:20]}...'\n\tenv_style: '{env_style[:20]}...'\n\tenv_negative: '{env_negative[:20]}...'")
        print(f"\twidth, height: ({width},{height})\n\tenv_cond_strength: {env_cond_strength}\n\tmat_cond_strength: {mat_cond_strength}\n\tipadapter_weight: {ipadapter_weight}")
        print(f"\tstart_at: {start_at}\n\tend_at: {end_at}")
        # Create a report string pairing each mat_txt with its corresponding mat_msk shape
        mat_report = "\n".join(
            f"\tmat[{i}]: {tuple(msk.shape)} '{txt[:20]}...'"
            for i, (txt, msk) in enumerate(zip(mat_txts, mat_msks))
        )
        print(f"{mat_report}")

        # Print relevant info about ipadapter contents
        if isinstance(ipadapter, dict):
            # Print clipvision info
            clipvision = ipadapter.get("clipvision", {})
            print(f"\tclipvision.model type: {type(clipvision.get('model', None))}")

            clipvision_file = clipvision.get('file', 'N/A')
            try: print(f"\tclipvision.file: {os.path.basename(clipvision_file)}")
            except Exception: print(f"\tclipvision.file: {clipvision_file}")

            # Print ipadapter file info
            ipadapter_dict = ipadapter.get("ipadapter", {})
            ipadapter_file = ipadapter_dict.get('file', 'N/A')
            try: print(f"\tipadapter.file: {os.path.basename(ipadapter_file)}")
            except Exception: print(f"\tipadapter.file: {ipadapter_file}")

            # Print keys in ipadapter.model
            ipadapter_model = ipadapter_dict.get("model", {})
            print(f"\tipadapter.model keys: {list(ipadapter_model.keys())}")


        # main
        #

        if start_at > end_at: 
            raise ValueError("start_at must be less than or equal to end_at.")

        # add styles to each object prompt
        styled_mat_txts = [ (env_style if m is None or m == '' else m + ", " + env_style) for m in mat_txts]
        
        # turn the list of strings into a list of conditionings:
        cond_mat_txts = [clip_text_encode(clip, s) for s in styled_mat_txts] # appending as a list - format of comfy when returning CONDITIONING type

        env_positive = "; ".join([env_scene, env_style]) # combine scene + style prompts into a single string

        positive_prompt_cond = clip_text_encode(clip, env_positive)
        negative_prompt_cond = clip_text_encode(clip, env_negative)

        empty_mask = create_solid_mask(1.0, width, height)
        positive_prompt_cond = conditioning_set_mask(positive_prompt_cond, empty_mask, strength=env_cond_strength) # first combining, using default vals for other params

        # apply text conditioning to positive_prompt_cond
        for i in range(len(cond_mat_txts)):
            temp = conditioning_set_mask( cond_mat_txts[i], mat_msks[i], strength=mat_cond_strength)
            positive_prompt_cond = conditioning_combine(positive_prompt_cond, temp)

        # apply image conditioning to the model
        for i in range(len(cond_mat_txts)):
            if mat_imgs[i] is not None:
                # model, ipadapter, image, weight, start_at, end_at, weight_type, attn_mask=None
                weight_type = 'standard' # 'standard' or 'style transfer' or 'prompt_is_more_important'
                model, _ = apply_ipadapter(model, ipadapter, mat_imgs[i], ipadapter_weight, start_at, end_at, weight_type, mat_msks[i])

        
        return (model, positive_prompt_cond, negative_prompt_cond)
    

# ==============================================================================
# utility functions
# ==============================================================================


def clip_text_encode(clip, str):
    tokens = clip.tokenize(str)
    cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
    return [[cond, {"pooled_output": pooled}]]

def conditioning_set_mask(conditioning, mask, set_cond_area="default", strength=1.0): # from builtin nodes: "append" func of the ConditioningSetMask node
        if not (0.0 <= strength <= 10.0):
            raise ValueError("Strength must be between 0.0 and 10.0.")

        set_area_to_bounds = False
        if set_cond_area != "default":
            set_area_to_bounds = True
        if len(mask.shape) < 3:
            mask = mask.unsqueeze(0)

        cond = node_helpers.conditioning_set_values(conditioning, {"mask": mask,
                                                                "set_area_to_bounds": set_area_to_bounds,
                                                                "mask_strength": strength})
        return cond

def conditioning_combine(conditioning_1, conditioning_2): # from builtin nodes: ConditioningCombine
    return conditioning_1 + conditioning_2