import torch
import logging
from typing import Optional

from comfy.language.transformers_model_management import TransformersManagedModel
from comfy.model_management import vram_state, VRAMState, unload_all_models, get_free_memory
from comfy.model_management_types import HooksSupport, ModelManageable
from comfy.model_patcher import ModelPatcher
from comfy.node_helpers import export_custom_nodes
from comfy.nodes.package_typing import CustomNode
from comfy.ops import manual_cast
from comfy.patcher_extension import WrappersMP
from comfy.rmsnorm import RMSNorm

_DISABLE_COMFYUI_CASTING_HOOK = "disable_comfyui_casting_hook"

logger = logging.getLogger(__name__)

_DIFFUSERS_IMPORT_ERROR: Optional[Exception] = None

try:
    from diffusers import HookRegistry
    from diffusers.hooks import apply_group_offloading, apply_layerwise_casting, ModelHook
except Exception as exc:  # pragma: no cover - depends on torch/diffusers runtime matrix
    _DIFFUSERS_IMPORT_ERROR = exc
    logger.warning("group offloading nodes disabled", exc_info=exc)


if _DIFFUSERS_IMPORT_ERROR is None:
    class DisableComfyWeightCast(ModelHook):
        r"""
        A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype
        for storage. This process may lead to quality loss in the output, but can significantly reduce the memory
        footprint.
        """

        _is_stateful = False

        def __init__(self) -> None:
            super().__init__()

        def initialize_hook(self, module: torch.nn.Module):
            if hasattr(module, "comfy_cast_weights"):
                module.comfy_cast_weights = False
            return module

        def deinitalize_hook(self, module: torch.nn.Module):
            if hasattr(module, "comfy_cast_weights"):
                module.comfy_cast_weights = True
            return module


    def disable_comfyui_weight_casting_hook(module: torch.nn.Module):
        registry = HookRegistry.check_if_exists_or_initialize(module)
        hook = DisableComfyWeightCast()
        registry.register_hook(hook, _DISABLE_COMFYUI_CASTING_HOOK)


    def disable_comfyui_weight_casting(module: torch.nn.Module):
        types = [
            torch.nn.Linear,
            torch.nn.Conv1d,
            torch.nn.Conv2d,
            torch.nn.Conv3d,
            torch.nn.GroupNorm,
            torch.nn.LayerNorm,
            RMSNorm,
            torch.nn.ConvTranspose2d,
            torch.nn.ConvTranspose1d,
            torch.nn.Embedding
        ]
        try:
            from torch.nn import RMSNorm as TorchRMSNorm
            types.append(TorchRMSNorm)
        except (ImportError, ModuleNotFoundError):
            pass

        if isinstance(module, tuple(types)):
            disable_comfyui_weight_casting_hook(module)
            return

        for _, submodule in module.named_children():
            disable_comfyui_weight_casting(submodule)


    def prepare_group_offloading_factory(load_device: torch.device, offload_device: torch.device):
        def wrapper(executor, model: ModelPatcher, *args, **kwargs):
            # this model will now just be loaded to CPU, since diffusers will manage moving to gpu
            model.load_device = offload_device

            # we'll have to unload everything to use pinning better, this includes trimming
            unload_all_models()

            # loads the model, prepares everything
            inner_model, conds, models = executor(model, *args, **kwargs)

            # we will need layer casting from diffusers in this situation
            if model.model.operations == manual_cast and model.diffusion_model.dtype != model.model.manual_cast_dtype:
                raise ValueError("manual casting operations, where the model is loaded in different weights than inference will occur, is not supported")

            # weights are patched, ready to go, inner model will be correctly deleted at the end of sampling
            model_size = model.model_size()

            model_too_large = model_size * 2 > get_free_memory(torch.cpu)
            low_vram_state = vram_state in (VRAMState.LOW_VRAM,)
            is_cuda_device = load_device.type == 'cuda'

            if model_too_large or low_vram_state:
                logger.error(f"group offloading did not use memory pinning because model_too_large={model_too_large} low_vram_state={low_vram_state}")
            if not is_cuda_device:
                logger.error(f"group offloading did not use stream because load_device.type={load_device.type} != \"cuda\"")
            apply_group_offloading(
                inner_model.diffusion_model,
                load_device,
                offload_device,
                use_stream=is_cuda_device,
                record_stream=is_cuda_device,
                low_cpu_mem_usage=low_vram_state or model_too_large,
                num_blocks_per_group=1
            )
            # then the inputs will be ready on the correct device due to the wrapper factory
            model.load_device = load_device
            return inner_model, conds, models

        return wrapper


    def prepare_layerwise_casting_factory(dtype: torch.dtype):
        def wrapper(executor, model: ModelPatcher, *args, **kwargs):
            disable_comfyui_weight_casting(model.diffusion_model)
            apply_layerwise_casting(model.diffusion_model,
                                    dtype,
                                    model.diffusion_model.dtype,
                                    non_blocking=True)
            inner_model, conds, models = executor(model, *args, **kwargs)

            return inner_model, conds, models

        return wrapper


    class GroupOffload(CustomNode):
        @classmethod
        def INPUT_TYPES(cls):
            return {
                "required": {
                    "model": ("MODEL", {})
                }
            }

        RETURN_TYPES = ("MODEL",)
        FUNCTION = "execute"

        def execute(self, model: ModelManageable | HooksSupport | TransformersManagedModel) -> tuple[ModelPatcher,]:
            if isinstance(model, ModelManageable):
                model = model.clone()
            if isinstance(model, TransformersManagedModel):
                apply_group_offloading(
                    model.model,
                    model.load_device,
                    model.offload_device,
                    use_stream=True,
                    record_stream=True,
                    low_cpu_mem_usage=vram_state in (VRAMState.LOW_VRAM,),
                    num_blocks_per_group=1
                )
            elif isinstance(model, HooksSupport) and isinstance(model, ModelManageable):
                model.add_wrapper_with_key(WrappersMP.PREPARE_SAMPLING, "group_offload", prepare_group_offloading_factory(model.load_device, model.offload_device))
            return model,


    class LayerwiseCast(CustomNode):
        @classmethod
        def INPUT_TYPES(cls):
            return {
                "required": {
                    "model": ("MODEL", {}),
                    "dtype": (["float8_e4m3fn", "float8_e5m2"], {})
                }
            }

        RETURN_TYPES = ("MODEL",)
        FUNCTION = "execute"

        def execute(self, model: ModelPatcher, dtype: str) -> tuple[ModelPatcher,]:
            model = model.clone()
            if dtype == "float8_e4m3fn":
                dtype = torch.float8_e4m3fn
            elif dtype == "float8_e5m2":
                dtype = torch.float8_e5m2

            model.add_wrapper(WrappersMP.PREPARE_SAMPLING, prepare_layerwise_casting_factory(dtype))
            return model,


export_custom_nodes()
