import torch
import torch.nn as nn
from safetensors.torch import safe_open
from .model import FluxParams, FluxMod
from .layers import Approximator
import comfy
import comfy.model_patcher
from comfy import model_management
import comfy.supported_models_base
import comfy.latent_formats
import comfy.model_patcher
import comfy.model_base
import comfy.utils
import comfy.conds
import comfy.ops


gguf = None

def ensure_gguf():
    global gguf
    if gguf:
        return
    import sys
    gguf = sys.modules.get("ComfyUI-GGUF")
    if gguf is None:
        raise RuntimeError("Could not find ComfyUI-GGUF node: GGUF support requires ComfyUI-GGUF")

class ExternalFlux(comfy.supported_models_base.BASE):
    unet_config = {}
    unet_extra_config = {}
    latent_format = comfy.latent_formats.Flux
    memory_usage_factor = 2.0 # Lower memory usage factor than FLUX, if we get OOMs this may need to be raised by ~0.2 or so.
      
    def __init__(self,):
        self.unet_config = {}
        self.latent_format = self.latent_format()
        self.unet_config["disable_unet_model_creation"] = True
    

class ExternalFluxModel(comfy.model_base.BaseModel):
    chroma_model_mode=False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def extra_conds(self, **kwargs):
        out = super().extra_conds(**kwargs)
        if self.chroma_model_mode:
            guidance = 0.0
        else:
            guidance = kwargs.get("guidance", 3.5)
            if guidance is None:
                guidance = 0.0
        out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor((guidance,)))
        return out

class ChromaFluxModel(ExternalFluxModel):
        chroma_model_mode=True

def load_selected_keys(filename, exclude_keywords=(), is_gguf=False):
    """Loads all keys from a safetensors file except those containing specified keywords.

    Args:
        filename: Path to the safetensors file.
        exclude_keywords: List of keywords to exclude.

    Returns:
        A dictionary containing the loaded tensors.
    """
    global gguf

    def is_excluded(key):
        return (
            not key.startswith("distilled_guidance_layer.") and
            any(keyword in key for keyword in exclude_keywords)
        )

    if is_gguf:
        # gguf_sd_loader will strip "model.diffusion_model." if it exists, so no need to handle it here.
        return {
            key: value
            for key, value in gguf.loader.gguf_sd_loader(filename).items()
            if not is_excluded(key)
        }
    tensors = {}
    if filename.endswith("pth"):
        state_dict = torch.load(filename, weights_only=True, map_location="cpu")
        for orig_key in state_dict.keys():
            if orig_key.startswith("model.diffusion_model."):
                key = orig_key[22:]
            else:
                key = orig_key
            if not is_excluded(key):
                tensors[key] = state_dict[orig_key]

    elif filename.endswith("safetensors"):
        with safe_open(filename, framework="pt") as f:
            for orig_key in f.keys():
                if orig_key.startswith("model.diffusion_model."):
                    key = orig_key[22:]
                else:
                    key = orig_key
                if not is_excluded(key):
                    tensors[key] = f.get_tensor(orig_key)
    else:
        raise NotImplementedError
    return tensors


def cast_layers(module, layer_type, dtype, exclude_keywords=()):
    """Casts layers in a module to the specified dtype, excluding layers with names containing keywords.

    Args:
        module (nn.Module): The module containing layers to cast.
        layer_type (type): The type of layers to cast (e.g., nn.Linear).
        dtype (torch.dtype): The target data type for casting.
        exclude_keywords (list, optional): A list of keywords to exclude from casting based on layer names. Defaults to [].

    Returns:
        None
    """

    for child_name, child in module.named_children():
        if isinstance(child, layer_type) and not any(keyword in child_name for keyword in exclude_keywords):
            # Cast the layer only if it's the desired type and not excluded
            # omit bias
            child.weight.data = child.weight.data.to(dtype=dtype)
        else:
            # Recursively apply to child modules
            cast_layers(child, layer_type, dtype, exclude_keywords)


def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.bfloat16, lite_patch_path=None, is_gguf=False):

    if is_gguf:
        ensure_gguf()
    # just load safetensors here
    state_dict = load_selected_keys(model_path, {"mod", "time_in", "guidance_in", "vector_in"}, is_gguf=is_gguf)

    if timestep_guidance_path is None:
        # Chroma mode - we expect the timestep guidance to be bundled under
        # the key distilled_guidance_layer.
        if lite_patch_path is not None:
            raise ValueError("Internal error: lite patch specified in Chroma loader mode")
        timestep_state_dict = {
            k[25:]: v
            for k, v in state_dict.items()
            if k.startswith("distilled_guidance_layer.")
        }
        if not timestep_state_dict:
            raise RuntimeError("Could not find distilled guidance layer in Chroma model")
        n_layers = 0
        for key in timestep_state_dict:
            keysplit = key.split(".", 2)
            if (
                len(keysplit) == 3 and
                keysplit[0] == "norms" and
                keysplit[1].isnumeric() and
                keysplit[2] == "scale"
            ):
                n_layers = max(n_layers, int(keysplit[1]) + 1)
            del state_dict[f"distilled_guidance_layer.{key}"]
        if n_layers == 0:
            raise RuntimeError("Could not determine number of distilled guidance layers in Chroma model")
    else:
        timestep_state_dict = comfy.utils.load_torch_file(timestep_guidance_path)

        if "v3" in timestep_guidance_path:
            n_layers = 6
        elif "v2" in timestep_guidance_path:
            n_layers = 5
        else:
            n_layers = 4

    param_count = sum(x.numel() for x in state_dict.values())

    load_device = model_management.get_torch_device()
    offload_device = model_management.unet_offload_device()
    unet_dtype = model_management.unet_dtype(model_params=param_count)
    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)


    params=FluxParams(
        in_channels=64,
        vec_in_dim=768,
        context_in_dim=4096,
        hidden_size=3072,
        mlp_ratio=4.0,
        num_heads=24,
        depth=19 if lite_patch_path is None else 8,
        depth_single_blocks=38,
        axes_dim=[16, 56, 56],
        theta=10_000,
        qkv_bias=True,
        guidance_embed=False,
    )

    model_conf = ExternalFlux()
    if state_dict["scaled_fp8"] is not None:
        scaled_fp8_weight = state_dict.pop("scaled_fp8")
        model_conf.scaled_fp8 = scaled_fp8_weight.dtype
        if model_conf.scaled_fp8 == torch.float32:
            model_conf.scaled_fp8 = torch.float8_e4m3fn
        if scaled_fp8_weight.nelement() == 2:
            model_conf.optimizations["fp8"] = False
        else:
            model_conf.optimizations["fp8"] = True
    model_class = ChromaFluxModel if timestep_guidance_path is None else ExternalFluxModel
    model = model_class(
        model_conf,
        model_type=comfy.model_base.ModelType.FLUX,
        device=load_device
    )
    unet_config = model_conf.unet_config
    if is_gguf:
        operations = model_conf.custom_operations = gguf.ops.GGMLOps()
        modelpatcher_class = gguf.nodes.GGUFModelPatcher
    else:
        modelpatcher_class = comfy.model_patcher.ModelPatcher
    if model_conf.custom_operations is None:
        fp8 = model_conf.optimizations.get("fp8", model_conf.scaled_fp8 is not None)
        operations = comfy.ops.pick_operations(
            unet_config.get("dtype"),
            model.manual_cast_dtype,
            fp8_optimizations=fp8,
            scaled_fp8=model_conf.scaled_fp8,
        )
    else:
        operations = model_conf.custom_operations
    
    model.diffusion_model = FluxMod(params=params, dtype=unet_dtype, operations=operations)

    model.diffusion_model.load_state_dict(state_dict)
    
    if lite_patch_path is not None:
        model.diffusion_model.lite = True
        # for _ in range(5,16):
        #     del model.diffusion_model.double_blocks[5]

        # model.diffusion_model.double_blocks[4].load_state_dict(comfy.utils.load_torch_file(lite_patch_path))

    model.diffusion_model.distilled_guidance_layer = Approximator(64 if lite_patch_path is None else 32, 3072, 5120, n_layers, operations=operations)
    model.diffusion_model.distilled_guidance_layer.load_state_dict(timestep_state_dict)
    model.diffusion_model.dtype = unet_dtype
    model.diffusion_model.eval()
    if not is_gguf:
        if model_conf.scaled_fp8 is None:
            #model.diffusion_model.to(linear_dtypes)
            # we cast to fp8 for mixed matmul ops but omit the picky and sensitive layers
            cast_layers(model.diffusion_model, nn.Linear, dtype=linear_dtypes, exclude_keywords={"img_in", "txt_in", "final_layer", "in_proj", "in_layer", "out_layer", "out_proj"})
        if model_management.force_channels_last():
            model.diffusion_model.to(memory_format=torch.channels_last)

    model_patcher = modelpatcher_class(
        model,
        load_device = load_device,
        offload_device = offload_device,
    )
    return model_patcher