# (c) City96 || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
import torch
import gguf
import logging
import numpy as np

import comfy.sd
import comfy.utils
import comfy.model_management
import folder_paths

from .ops import GGMLTensor, GGMLOps
from .lora import load_lora_gguf

# TODO: This causes gguf files to show up in the main unet loader
folder_paths.folder_names_and_paths["unet"][1].add(".gguf")

def gguf_sd_loader(path):
    """
    Read state dict as fake tensors
    """
    reader = gguf.GGUFReader(path)
    sd = {}
    dt = {}
    for tensor in reader.tensors:
        sd[str(tensor.name)] = GGMLTensor(
            torch.from_numpy(tensor.data), # mmap
            tensor_type = tensor.tensor_type,
            tensor_shape = torch.Size(
                np.flip(list(tensor.shape))
            )
        )
        dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1

    # sanity check debug print
    print("\nggml_sd_loader:")
    for k,v in dt.items():
        print(f" {k:30}{v:3}")
    print("\n")
    return sd

class UnetLoaderGGUF:
    @classmethod
    def INPUT_TYPES(s):
        unet_names = [x for x in folder_paths.get_filename_list("unet") if x.endswith(".gguf")]
        return {
            "required": {
                "unet_name": (unet_names,),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_unet"
    CATEGORY = "bootleg"
    TITLE = "Unet Loader (GGUF)"

    def load_unet(self, unet_name):
        unet_path = folder_paths.get_full_path("unet", unet_name)
        sd = gguf_sd_loader(unet_path)
        model = comfy.sd.load_diffusion_model_state_dict(
            sd, model_options={"custom_operations": GGMLOps}
        )
        if model is None:
            logging.error("ERROR UNSUPPORTED UNET {}".format(unet_path))
            raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
        return (model,)

class LoraLoaderGGUFModelOnly:
    def __init__(self):
        self.loaded_lora = None

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                 "model": ("MODEL",),
                 "lora_name": (folder_paths.get_filename_list("loras"), ),
                 "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
                 "lora_on_gpu": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "load_lora"
    CATEGORY = "bootleg"
    TITLE = "LoRA Loader Model Only (GGUF)"

    def load_lora(self, model, lora_name, strength_model, lora_on_gpu):
        if strength_model == 0:
            return (model,)

        lora_path = folder_paths.get_full_path("loras", lora_name)
        lora = None
        if self.loaded_lora is not None:
            if self.loaded_lora[0] == lora_path:
                lora = self.loaded_lora[1]
            else:
                temp = self.loaded_lora
                self.loaded_lora = None
                del temp

        if lora is None:
            lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
            self.loaded_lora = (lora_path, lora)

        # This is a super hacky method to get it to load alongside the model
        param_dict_name = f"lora_state_dict"#_{lora_name}".replace('.', '_')
        target_model = model.model.diffusion_model
        if lora_on_gpu:
            print(f"WARNING! Loading the LoRA to GPU memory may cause OOM issues or memory leaks!")
            temp = {}
            params = torch.nn.ParameterDict()
            for k,v in lora.items():
                param = torch.nn.Parameter(v, requires_grad=False)
                temp[k] = param
                params[k.replace('.', '_')] = param
            params.to(model.load_device, dtype=torch.float16)
            target_model.register_module(param_dict_name, params)
            lora = temp
        else:
            params = getattr(target_model, param_dict_name, None)
            if params:
                params.to("cpu", dtype=torch.float16)
                delattr(target_model, param_dict_name)

        model_lora = load_lora_gguf(model, lora, strength_model)
        return (model_lora,)

NODE_CLASS_MAPPINGS = {
    "UnetLoaderGGUF": UnetLoaderGGUF,
    "LoraLoaderGGUFModelOnly": LoraLoaderGGUFModelOnly,
}