import os
from collections import OrderedDict

import torch

from backend import memory_management
from backend.nn.cnets import cldm
from backend.operations import using_forge_operations
from backend.patcher.controlnet import (
    ControlLora,
    ControlNet,
    apply_controlnet_advanced,
    load_t2i_adapter,
    logger,
)
from modules import shared
from modules_forge.packages.comfy.utils import unet_to_diffusers
from modules_forge.packages.huggingface_guess.detection import (
    model_config_from_unet,
    unet_config_from_diffusers_unet,
)
from modules_forge.shared import add_supported_control_model


class ModelCache:
    def __init__(self):
        self.cache = OrderedDict()

    def get(self, key):
        if key in self.cache:
            self.cache.move_to_end(key)
            return self.cache[key]
        return None

    def put(self, key, model):
        self.cache[key] = model
        self.cache.move_to_end(key)

        max_size = max(1, shared.opts.data.get("control_net_model_cache_size", 3))
        while len(self.cache) > max_size:
            self.cache.popitem(last=False)


_CONTROL_MODEL_CACHE = ModelCache()


class ControlModelPatcher:

    @staticmethod
    def try_build_from_state_dict(state_dict, ckpt_path):
        return None

    def __init__(self, model_patcher=None):
        self.model_patcher = model_patcher
        self.strength = 1.0
        self.start_percent = 0.0
        self.end_percent = 1.0
        self.positive_advanced_weighting = None
        self.negative_advanced_weighting = None
        self.advanced_frame_weighting = None
        self.advanced_sigma_weighting = None
        self.advanced_mask_weighting = None

    def process_after_running_preprocessors(self, process, params, *args, **kwargs):
        return

    def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):
        return

    def process_after_every_sampling(self, process, params, *args, **kwargs):
        return


class ControlNetPatcher(ControlModelPatcher):

    def __init__(self, model_patcher):
        super().__init__(model_patcher)

    @staticmethod
    def try_build_from_state_dict(controlnet_data: dict[str, torch.Tensor], ckpt_path):
        if "lora_controlnet" in controlnet_data:
            return ControlNetPatcher(ControlLora(controlnet_data))

        controlnet_config = None
        if "controlnet_cond_embedding.conv_in.weight" in controlnet_data:  # diffusers format
            unet_dtype = memory_management.unet_dtype()
            controlnet_config = unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
            diffusers_keys = unet_to_diffusers(controlnet_config)
            diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
            diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"

            count = 0
            loop = True
            while loop:
                suffix = [".weight", ".bias"]
                for s in suffix:
                    k_in = "controlnet_down_blocks.{}{}".format(count, s)
                    k_out = "zero_convs.{}.0{}".format(count, s)
                    if k_in not in controlnet_data:
                        loop = False
                        break
                    diffusers_keys[k_in] = k_out
                count += 1

            count = 0
            loop = True
            while loop:
                suffix = [".weight", ".bias"]
                for s in suffix:
                    if count == 0:
                        k_in = "controlnet_cond_embedding.conv_in{}".format(s)
                    else:
                        k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
                    k_out = "input_hint_block.{}{}".format(count * 2, s)
                    if k_in not in controlnet_data:
                        k_in = "controlnet_cond_embedding.conv_out{}".format(s)
                        loop = False
                    diffusers_keys[k_in] = k_out
                count += 1

            new_sd = {}
            for k in diffusers_keys:
                if k in controlnet_data:
                    new_sd[diffusers_keys[k]] = controlnet_data.pop(k)

            if "control_add_embedding.linear_1.bias" in controlnet_data:  # Union Controlnet
                controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0]
                for k in list(controlnet_data.keys()):
                    new_k = k.replace(".attn.in_proj_", ".attn.in_proj.")
                    new_sd[new_k] = controlnet_data.pop(k)

            leftover_keys = controlnet_data.keys()
            if len(leftover_keys) > 0:
                logger.warning("Leftover Keys: {}".format(leftover_keys))
            controlnet_data = new_sd

        pth_key = "control_model.zero_convs.0.0.weight"
        pth = False
        key = "zero_convs.0.0.weight"
        if pth_key in controlnet_data:
            pth = True
            key = pth_key
            prefix = "control_model."
        elif key in controlnet_data:
            prefix = ""
        else:
            net = load_t2i_adapter(controlnet_data)
            if net is None:
                if not any(k.startswith("lllite") for k in controlnet_data):  # LLLite
                    logger.error("Could not detect Control model type...")
                return None
            return ControlNetPatcher(net)

        if controlnet_config is None:
            unet_dtype = memory_management.unet_dtype()
            controlnet_config = model_config_from_unet(controlnet_data, prefix, True).unet_config
            controlnet_config["dtype"] = unet_dtype

        load_device = memory_management.get_torch_device()
        computation_dtype = shared.sd_model.forge_objects.unet.model.computation_dtype
        controlnet_config.pop("out_channels")
        controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
        controlnet_config["hint_width"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[0]

        global _CONTROL_MODEL_CACHE
        control_model = _CONTROL_MODEL_CACHE.get(ckpt_path)

        if control_model is not None:
            logger.info("Reusing ControlNet Model...")
        else:
            logger.info("Creating ControlNet Model...")
            with using_forge_operations(dtype=unet_dtype, manual_cast_enabled=computation_dtype != unet_dtype):
                control_model = cldm.ControlNet(**controlnet_config).to(dtype=unet_dtype)

            if pth:
                if "difference" in controlnet_data:
                    logger.warning("Please use an official Control model for better performance...")

                class WeightsLoader(torch.nn.Module):
                    pass

                w = WeightsLoader()
                w.control_model = control_model
                missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
            else:
                missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)

            if len(missing) > 0:
                logger.warning("Missing ControlNet Keys: {}".format(missing))

            if len(unexpected) > 0:
                logger.debug("Unexpected ControlNet Keys: {}".format(unexpected))

            _CONTROL_MODEL_CACHE.put(ckpt_path, control_model)

        global_average_pooling = False
        filename = os.path.splitext(ckpt_path)[0]
        if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"):
            # TODO: smarter way of enabling global_average_pooling
            global_average_pooling = True

        control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=computation_dtype)
        return ControlNetPatcher(control)

    def process_before_every_sampling(self, process, cond, mask, *args, control_type=None, **kwargs):
        unet = process.sd_model.forge_objects.unet
        unet = apply_controlnet_advanced(
            unet=unet,
            controlnet=self.model_patcher,
            image_bchw=cond,
            strength=self.strength,
            start_percent=self.start_percent,
            end_percent=self.end_percent,
            positive_advanced_weighting=self.positive_advanced_weighting,
            negative_advanced_weighting=self.negative_advanced_weighting,
            advanced_frame_weighting=self.advanced_frame_weighting,
            advanced_sigma_weighting=self.advanced_sigma_weighting,
            advanced_mask_weighting=self.advanced_mask_weighting,
            control_type=control_type,
        )
        process.sd_model.forge_objects.unet = unet


add_supported_control_model(ControlNetPatcher)
