import folder_paths
import comfy.controlnet
import comfy.model_management
from nodes import NODE_CLASS_MAPPINGS

union_controlnet_types = {"auto": -1, "openpose": 0, "depth": 1, "hed/pidi/scribble/ted": 2, "canny/lineart/anime_lineart/mlsd": 3, "normal": 4, "segment": 5, "tile": 6, "repaint": 7}

class easyControlnet:
    def __init__(self):
        pass

    def apply(self, control_net_name, image, positive, negative, strength, start_percent=0, end_percent=1, control_net=None, scale_soft_weights=1, mask=None, union_type=None, easyCache=None, use_cache=True, model=None, vae=None):
        if strength == 0:
            return (positive, negative)

        # kolors controlnet patch
        from ..kolors.loader import is_kolors_model, applyKolorsUnet
        if is_kolors_model(model):
            from ..kolors.model_patch import patch_controlnet
            if control_net is None:
                with applyKolorsUnet():
                    control_net = easyCache.load_controlnet(control_net_name, scale_soft_weights, use_cache)
                control_net = patch_controlnet(model, control_net)
        else:
            if control_net is None:
                control_net = easyCache.load_controlnet(control_net_name, scale_soft_weights, use_cache)

        # union controlnet
        if union_type is not None:
            control_net = control_net.copy()
            type_number = union_controlnet_types[union_type]
            if type_number >= 0:
                control_net.set_extra_arg("control_type", [type_number])
            else:
                control_net.set_extra_arg("control_type", [])

        if mask is not None:
            mask = mask.to(self.device)

        if mask is not None and len(mask.shape) < 3:
            mask = mask.unsqueeze(0)

        control_hint = image.movedim(-1, 1)

        is_cond = True
        if negative is None:
            p = []
            for t in positive:
                n = [t[0], t[1].copy()]
                c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent))
                if 'control' in t[1]:
                    c_net.set_previous_controlnet(t[1]['control'])
                n[1]['control'] = c_net
                n[1]['control_apply_to_uncond'] = True
                if mask is not None:
                    n[1]['mask'] = mask
                    n[1]['set_area_to_bounds'] = False
                p.append(n)
            positive = p
        else:
            cnets = {}
            out = []
            for conditioning in [positive, negative]:
                c = []
                for t in conditioning:
                    d = t[1].copy()

                    prev_cnet = d.get('control', None)
                    if prev_cnet in cnets:
                        c_net = cnets[prev_cnet]
                    else:
                        c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae)
                        c_net.set_previous_controlnet(prev_cnet)
                        cnets[prev_cnet] = c_net

                    d['control'] = c_net
                    d['control_apply_to_uncond'] = False

                    if mask is not None:
                        d['mask'] = mask
                        d['set_area_to_bounds'] = False

                    n = [t[0], d]
                    c.append(n)
                out.append(c)
            positive = out[0]
            negative = out[1]

        return (positive, negative)