import os
import torch
import logging
from contextlib import contextmanager

# Model dep
import torch

try:
    import triton

    use_triton = True
except ImportError:
    use_triton = False

## XUT is the custom arch used in HDM
from .xut import env
from comfy.cli_args import args

env.TORCH_COMPILE = use_triton and not args.cpu
env.USE_LIGER = False
env.USE_XFORMERS = False
env.USE_XFORMERS_LAYERS = True
from .xut.xut import XUDiT
from .xut.modules.axial_rope import make_axial_pos_no_cache
from transformers import Qwen3Model, Qwen2Tokenizer

# Comfy
import folder_paths
import comfy.model_management
import comfy.latent_formats as latent_formats
import comfy.sd1_clip as sd1_clip
import comfy.sd
import comfy.model_management as model_management
import comfy.model_patcher
import comfy.utils
import comfy.conds
from comfy.supported_models_base import BASE, ClipTarget
from comfy.model_base import BaseModel, ModelType


@contextmanager
def operation_patch(operations):
    """
    Directly patch torch.nn so we don't need to re-impl the model arch
    Not recommended but easiest
    """
    module_list = set(i for i in dir(torch.nn) if not i.startswith("_"))
    for op in dir(operations):
        if op in module_list:
            setattr(torch.nn, f"org_{op}", getattr(torch.nn, op))
            setattr(torch.nn, op, getattr(operations, op))
    yield
    for op in dir(operations):
        if op in module_list:
            setattr(torch.nn, op, getattr(torch.nn, f"org_{op}"))


class EQVAE(latent_formats.LatentFormat):
    """
    Latent Format for https://huggingface.co/KBlueLeaf/EQ-SDXL-VAE
    """

    def __init__(self):
        self.scale_factor = 1.0
        self.latents_mean = torch.tensor(
            [
                -3.865960634760459,
                -1.3353925720045965,
                0.8353661802812418,
                0.909888159310023,
            ]
        ).view(1, 4, 1, 1)
        self.latents_std = torch.tensor(
            [
                10.036514488374861,
                6.848430044716239,
                7.213817552213277,
                5.591315799887989,
            ]
        ).view(1, 4, 1, 1)

    def process_in(self, latent):
        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
        latents_std = self.latents_std.to(latent.device, latent.dtype)
        return (latent - latents_mean) * self.scale_factor / latents_std

    def process_out(self, latent):
        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
        latents_std = self.latents_std.to(latent.device, latent.dtype)
        return latent * latents_std / self.scale_factor + latents_mean


class Qwen3_600MTokenizer(sd1_clip.SDTokenizer):
    """
    Tokenizer Class for Qwen3 with TI embedding support
    """

    def __init__(self, embedding_directory=None, tokenizer_data={}):
        tokenizer_path = os.path.join(
            os.path.dirname(os.path.realpath(__file__)), "qwen3_tokenizer"
        )
        super().__init__(
            tokenizer_path,
            pad_with_end=False,
            embedding_size=1024,  # We need this for TI embedding
            embedding_key="qwen3",
            tokenizer_class=Qwen2Tokenizer,
            has_start_token=False,
            has_end_token=False,
            pad_to_max_length=False,
            max_length=99999999,
            min_length=1,
            pad_token=151643,
            tokenizer_data=tokenizer_data,
        )


class Qwen3_600M_Wrapper(Qwen3Model):
    """
    Qwen3 Model wrapper to support ComfyUI's API
    """

    def forward(
        self,
        input_tokens,
        attention_mask=None,
        embeds=None,
        num_tokens=None,
        intermediate_output=None,
        final_layer_norm_intermediate=True,
        dtype=None,
        **kwargs,
    ):
        # We use Autocast here to resolve dtype issue easily
        # In theory you may want to manually set dtype instead of autocast for performance
        with torch.autocast(self.device.type, dtype=self.dtype):
            result = (
                super()
                .forward(input_tokens, attention_mask, inputs_embeds=embeds)
                .last_hidden_state
            )
        # we only have last layer output, and we don't have pooled emb as well
        return result.to(dtype), None


def Qwen3_600M(config, dtype, device, operations):
    """
    Qwen3 Model Class, we init model in normal mode + operation patch
    num_layers is needed in ComfyUI's Node
    """
    with torch.inference_mode(False), operation_patch(operations):
        model = (
            Qwen3_600M_Wrapper.from_pretrained(
                "Qwen/Qwen3-0.6B", torch_dtype=dtype, attn_implementation="sdpa"
            )
            .to(device)
            .to(dtype)
        )
        model = model.to(device).eval().requires_grad_(False)
        model.num_layers = model.config.num_hidden_layers
    return model


class Qwen3_600MModel(sd1_clip.SDClipModel):
    """
    The TextEncoder class for ComfyUI
    if you have your own following API, you don't need to inheirt this class:
        process_tokens(), load_sd(), encode(), reset/set_clip_options()
    Just make a child class of sd1_clip.SD1ClipModel will be easier
    """

    def __init__(
        self,
        device="cpu",
        layer="last",
        layer_idx=None,
        dtype=None,
        attention_mask=True,
        model_options={},
    ):
        super().__init__(
            device=device,
            layer=layer,
            layer_idx=layer_idx,
            textmodel_json_config={},
            dtype=dtype,
            special_tokens={"pad": 151643},
            layer_norm_hidden_state=False,
            model_class=Qwen3_600M,
            enable_attention_masks=attention_mask,
            return_attention_masks=attention_mask,
            model_options=model_options,
        )


class HDMQwenModel(sd1_clip.SD1ClipModel):
    """
    TE(CLIP) API wrapper in ComfyUI
    """

    def __init__(self, device="cpu", dtype=None, model_options={}):
        super().__init__(
            device=device,
            dtype=dtype,
            name="qwen3_600m",  # This name will be used to find TE in state_dict
            clip_model=Qwen3_600MModel,
            model_options=model_options,
        )

    def encode_token_weights(self, token_weight_pairs):
        result = self.qwen3_600m.encode_token_weights(token_weight_pairs)
        return result


class HDMWrapper(XUDiT):
    """
    Wrapper for initialize XUDiT model and do forward in ComfyUI
    """

    def __init__(
        self,
        patch_size=2,
        input_dim=4,
        dim=1024,
        ctx_dim=1024,
        ctx_size=256,
        heads=16,
        dim_head=64,
        mlp_dim=3072,
        depth=4,
        enc_blocks=1,
        mid_blocks=0,
        dec_blocks=3,
        dec_ctx=False,
        class_cond=0,
        shared_adaln=True,
        concat_ctx=True,
        use_dyt=False,
        self_ctx_mode="xut",
        double_t=False,
        addon_info_embs_dim=1,
        tread_config={
            "prev_trns_depth": 1,
            "post_trns_depth": 3,
            "dropout_ratio": 0.5,
        },
        dtype=None,
        device=None,
        operations=None,
        **kwargs,
    ):
        with torch.inference_mode(False), operation_patch(operations):
            super().__init__(
                patch_size=patch_size,
                input_dim=input_dim,
                dim=dim,
                ctx_dim=ctx_dim,
                ctx_size=ctx_size,
                heads=heads,
                dim_head=dim_head,
                mlp_dim=mlp_dim,
                depth=depth,
                enc_blocks=enc_blocks,
                mid_blocks=mid_blocks,
                dec_blocks=dec_blocks,
                dec_ctx=dec_ctx,
                class_cond=class_cond,
                shared_adaln=shared_adaln,
                concat_ctx=concat_ctx,
                use_dyt=use_dyt,
                self_ctx_mode=self_ctx_mode,
                double_t=double_t,
                addon_info_embs_dim=addon_info_embs_dim,
                tread_config=tread_config,
            )
            self.eval().requires_grad_(False)
            if device is not None:
                self.to(device)
            if dtype is not None:
                self.to(dtype)
            self.dtype = dtype
            self.device = device

    def forward(
        self,
        x,
        timesteps,
        context,
        x_shift=None,
        y_shift=None,
        zoom=None,
        tread_gamma=None,
        *args,
        **kwargs,
    ):
        b, c, h, w = x.shape
        aspect_ratio = w / h
        aspect_ratio_info = (
            torch.tensor([aspect_ratio], dtype=x.dtype, device=x.device)
            .repeat(b, 1)
            .log()
        )

        if x_shift is not None:
            pos_map = make_axial_pos_no_cache(h, w, dtype=x.dtype, device=x.device)[
                None
            ].repeat(b, 1, 1)
            if x_shift.ndim == 1:
                x_shift = x_shift[:, None]
                y_shift = y_shift[:, None]
                zoom = zoom[:, None, None]
            pos_map[..., 1] = pos_map[..., 1] + x_shift
            pos_map[..., 0] = pos_map[..., 0] + y_shift
            pos_map = pos_map / zoom
        else:
            pos_map = None

        tread_gamma = tread_gamma[0].item() if tread_gamma is not None else None
        tread_gamma = tread_gamma or None

        return super().forward(
            x.to(self.dtype),
            timesteps.to(self.dtype),
            context.to(self.dtype),
            pos_map=pos_map,
            addon_info=aspect_ratio_info,
            tread_rate=tread_gamma,
        )


class HDMSmall(BaseModel):
    def __init__(self, model_config, model_type=ModelType.FLOW, device=None, **kwargs):
        model_config.unet_config.update(kwargs)
        super().__init__(model_config, model_type, device=device, unet_model=HDMWrapper)

    def extra_conds(self, **kwargs):
        base_result = super().extra_conds(**kwargs)
        if "x_shift" in kwargs:
            x_shift, y_shift, zoom = (
                kwargs["x_shift"],
                kwargs["y_shift"],
                kwargs["zoom"],
            )
            x_shift = torch.tensor([x_shift], dtype=torch.float32, device=self.device)
            y_shift = torch.tensor([y_shift], dtype=torch.float32, device=self.device)
            zoom = torch.tensor([zoom], dtype=torch.float32, device=self.device)
            base_result["x_shift"] = comfy.conds.CONDRegular(x_shift)
            base_result["y_shift"] = comfy.conds.CONDRegular(y_shift)
            base_result["zoom"] = comfy.conds.CONDRegular(zoom)
        if "tread_gamma" in kwargs:
            tread_gamma = torch.tensor(
                [kwargs["tread_gamma"]], dtype=torch.float32, device=self.device
            )
            base_result["tread_gamma"] = comfy.conds.CONDRegular(tread_gamma)
        return base_result


class HomeDiffusionSmall(BASE):
    unet_config = {}

    # Standard flow matching without shift
    sampling_settings = {
        "multiplier": 1.0,
        "shift": 1,
    }

    # Some random factor here, should be determined by profiling
    memory_usage_factor = 2

    unet_extra_config = {}
    latent_format = EQVAE

    # HDM is trained with fp16-amp so all dtype supported
    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]

    vae_key_prefix = ["vae."]
    text_encoder_key_prefix = ["te.text_models.0."]

    def __init__(self, unet_config):
        super().__init__(unet_config)
        if comfy.model_management.extended_fp16_support():
            self.supported_inference_dtypes = [
                torch.float16
            ] + self.supported_inference_dtypes

    def get_model(self, state_dict, prefix="", device=None):
        max_num_mid_blocks = 0
        for key in list(state_dict.keys()):
            if "mid_blocks." in key:
                number = int(key.split("mid_blocks.")[1].split(".")[0]) + 1
                max_num_mid_blocks = max(max_num_mid_blocks, number)
        out = HDMSmall(self, mid_blocks=max_num_mid_blocks, device=device)
        return out

    def clip_target(self, state_dict={}):
        return ClipTarget(Qwen3_600MTokenizer, HDMQwenModel)

    def process_clip_state_dict(self, state_dict):
        """
        Custom state dict process for loading TE
        """
        state_dict = comfy.utils.state_dict_prefix_replace(
            state_dict,
            {"te.text_models.0.": "qwen3_600m.transformer."},
            filter_keys=True,
        )
        state_dict["qwen3_600m.logit_scale"] = torch.tensor(1.0)
        return state_dict


def load_state_dict_hdm(
    sd,
    embedding_directory=None,
    model_options={},
    te_model_options={},
    metadata=None,
):
    """Directly copied from comfy.sd and simplified for HDM only"""
    clip = None
    vae = None
    model = None
    model_patcher = None
    diffusion_model_prefix = "unet.model."

    parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
    weight_dtype = comfy.utils.weight_dtype(sd, diffusion_model_prefix)
    load_device = model_management.get_torch_device()

    model_config = HomeDiffusionSmall({})

    unet_weight_dtype = list(model_config.supported_inference_dtypes)
    if model_config.scaled_fp8 is not None:
        weight_dtype = None

    model_config.custom_operations = model_options.get("custom_operations", None)
    unet_dtype = model_options.get("dtype", model_options.get("weight_dtype", None))

    if unet_dtype is None:
        unet_dtype = model_management.unet_dtype(
            model_params=parameters,
            supported_dtypes=unet_weight_dtype,
            weight_dtype=weight_dtype,
        )

    manual_cast_dtype = model_management.unet_manual_cast(
        unet_dtype, load_device, model_config.supported_inference_dtypes
    )
    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)

    inital_load_device = model_management.unet_inital_load_device(
        parameters, unet_dtype
    )
    model = model_config.get_model(
        sd, diffusion_model_prefix, device=inital_load_device
    )
    model.load_model_weights(sd, diffusion_model_prefix)

    # In HDM we use SDXL vae arch so we can directly use built-in VAE class here
    # For custom VAE you will need to implement encode() and decode() method
    vae_sd = comfy.utils.state_dict_prefix_replace(
        sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True
    )
    vae_sd = model_config.process_vae_state_dict(vae_sd)
    vae = comfy.sd.VAE(sd=vae_sd, metadata=metadata)

    clip_target = model_config.clip_target(state_dict=sd)
    if clip_target is not None:
        clip_sd = model_config.process_clip_state_dict(sd)
        parameters = comfy.utils.calculate_parameters(clip_sd)
        if parameters < 1e6:
            # no TE in ckpt, use pre-defined value:
            # 600M
            parameters = 600_000_000
        clip = comfy.sd.CLIP(
            clip_target,
            embedding_directory=embedding_directory,
            tokenizer_data=clip_sd,
            parameters=parameters,
            model_options=te_model_options,
        )
        m, u = clip.load_sd(clip_sd, full_model=True)
        if len(m) > 0:
            logging.warning(
                "HDM TE missing some weights in state dict, "
                "if you are using DiT only ckpt than this is expected."
            )

        if len(u) > 0:
            logging.debug("clip unexpected {}:".format(u))

    left_over = sd.keys()
    if len(left_over) > 0:
        logging.debug("left over keys: {}".format(left_over))

    model_patcher = comfy.model_patcher.ModelPatcher(
        model,
        load_device=load_device,
        offload_device=model_management.unet_offload_device(),
    )
    if inital_load_device != torch.device("cpu"):
        logging.info("loaded diffusion model directly to GPU")
        model_management.load_models_gpu([model_patcher], force_full_load=True)

    return (model_patcher, clip, vae)


def load_checkpoint_hdm(
    ckpt_path,
    embedding_directory=None,
    model_options={},
    te_model_options={},
):
    sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
    out = load_state_dict_hdm(
        sd,
        embedding_directory,
        model_options,
        te_model_options=te_model_options,
        metadata=metadata,
    )
    if out is None:
        raise RuntimeError("ERROR: Could not load HDM model")
    return out


class HDMCheckpointLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ckpt_name": (
                    folder_paths.get_filename_list("checkpoints"),
                    {"tooltip": "The name of the checkpoint (model) to load."},
                ),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    OUTPUT_TOOLTIPS = (
        "The model used for denoising latents.",
        "The CLIP model used for encoding text prompts.",
        "The VAE model used for encoding and decoding images to and from latent space.",
    )
    FUNCTION = "load_checkpoint"

    CATEGORY = "loaders"
    DESCRIPTION = "Loads a HDM model checkpoint."

    def load_checkpoint(self, ckpt_name):
        ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
        return load_checkpoint_hdm(
            ckpt_path,
            embedding_directory=folder_paths.get_folder_paths("embeddings"),
        )


class HDMCameraParam:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "x_shift": (
                    "FLOAT",
                    {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
                ),
                "y_shift": (
                    "FLOAT",
                    {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01},
                ),
                "zoom": (
                    "FLOAT",
                    {"default": 1.0, "min": 0.2, "max": 5.0, "step": 0.01},
                ),
            }
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
    OUTPUT_TOOLTIPS = ""
    FUNCTION = "apply"

    CATEGORY = "conditioning/hdm"
    DESCRIPTION = "Apply camera parameters to conditioning"

    def apply(self, positive, negative, x_shift, y_shift, zoom):
        for pos in positive:
            pos[1]["x_shift"] = x_shift
            pos[1]["y_shift"] = y_shift
            pos[1]["zoom"] = zoom
        for neg in negative:
            neg[1]["x_shift"] = x_shift
            neg[1]["y_shift"] = y_shift
            neg[1]["zoom"] = zoom
        return positive, negative


class HDMTreadGamma:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "positive_gamma": (
                    "FLOAT",
                    {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
                "negative_gamma": (
                    "FLOAT",
                    {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01},
                ),
            }
        }

    RETURN_TYPES = ("CONDITIONING", "CONDITIONING")
    OUTPUT_TOOLTIPS = ""
    FUNCTION = "apply"

    CATEGORY = "conditioning/hdm"
    DESCRIPTION = "Apply tread gamma to conditioning"

    def apply(self, positive, negative, positive_gamma, negative_gamma):
        for pos in positive:
            pos[1]["tread_gamma"] = positive_gamma
        for neg in negative:
            neg[1]["tread_gamma"] = negative_gamma
        return positive, negative


NODE_CLASS_MAPPINGS = {
    "HDMLoader": HDMCheckpointLoader,
    "HDMCameraParam": HDMCameraParam,
    "HDMTreadGamma": HDMTreadGamma,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "HDMLoader": "HDM Loader",
    "HDMCameraParam": "HDM Camera",
    "HDMTreadGamma": "HDM Tread Gamma",
}