import torch
from ..utils import log
import comfy.model_management as mm

device = mm.get_torch_device()
offload_device = mm.unet_offload_device()

class WanVideoAddSCAILReferenceEmbeds:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "embeds": ("WANVIDIMAGE_EMBEDS",),
                    "vae": ("WANVAE", {"tooltip": "VAE model"}),
                    "ref_image": ("IMAGE",),
                    "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the reference embedding"}),
                    "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the embedding application"}),
                    "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the embedding application"}),
                },
                "optional": {
                    "clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
                }
        }

    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
    RETURN_NAMES = ("image_embeds",)
    FUNCTION = "add"
    CATEGORY = "WanVideoWrapper"

    def add(self, embeds, vae, ref_image, strength, start_percent, end_percent, clip_embeds=None):
        updated = dict(embeds)

        vae.to(device)
        ref_image_in = (ref_image[..., :3].permute(3, 0, 1, 2) * 2 - 1).to(device, vae.dtype)
        ref_latent = vae.encode([ref_image_in], device, tiled=False)[0]
        log.info(f"SCAIL ref_latent shape: {ref_latent.shape}")

        ref_mask = torch.ones_like(ref_latent[:4])
        ref_latent = torch.cat([ref_latent, ref_mask], dim=0)
        vae.to(offload_device)

        updated.setdefault("scail_embeds", {})
        updated["scail_embeds"]["ref_latent_pos"] = ref_latent * strength
        updated["scail_embeds"]["ref_latent_neg"] = torch.zeros_like(ref_latent)
        updated["scail_embeds"]["ref_start_percent"] = start_percent
        updated["scail_embeds"]["ref_end_percent"] = end_percent
        updated["clip_context"] = clip_embeds.get("clip_embeds", None) if clip_embeds is not None else None

        return (updated,)

class WanVideoAddSCAILPoseEmbeds:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
                    "embeds": ("WANVIDIMAGE_EMBEDS",),
                    "vae": ("WANVAE", {"tooltip": "VAE model"}),
                    "pose_images": ("IMAGE", {"tooltip": "Pose images for the entire video"}),
                    "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Strength of the pose control"}),
                    "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percentage of the pose control application"}),
                    "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percentage of the pose control application"}),
                },
        }

    RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
    RETURN_NAMES = ("image_embeds",)
    FUNCTION = "add"
    CATEGORY = "WanVideoWrapper"

    def add(self, embeds, vae, pose_images, strength, start_percent=0.0, end_percent=1.0):
        updated = dict(embeds)

        vae.to(device)
        pose_images_in = (pose_images[..., :3].permute(3, 0, 1, 2) * 2 - 1).to(device, vae.dtype)
        pose_latent = vae.encode([pose_images_in], device, tiled=False)[0]
        pose_mask = torch.ones_like(pose_latent[:4])
        pose_latent = torch.cat([pose_latent, pose_mask], dim=0)
        log.info(f"SCAIL pose_latent shape: {pose_latent.shape}")

        vae.to(offload_device)

        updated.setdefault("scail_embeds", {})
        updated["scail_embeds"]["pose_latent"] = pose_latent
        updated["scail_embeds"]["pose_strength"] = strength
        updated["scail_embeds"]["pose_start_percent"] = start_percent
        updated["scail_embeds"]["pose_end_percent"] = end_percent

        return (updated,)


NODE_CLASS_MAPPINGS = {
    "WanVideoAddSCAILPoseEmbeds": WanVideoAddSCAILPoseEmbeds,
    "WanVideoAddSCAILReferenceEmbeds": WanVideoAddSCAILReferenceEmbeds,
    }
NODE_DISPLAY_NAME_MAPPINGS = {
    "WanVideoAddSCAILReferenceEmbeds": "WanVideo Add SCAIL Reference Embeds",
    "WanVideoAddSCAILPoseEmbeds": "WanVideo Add SCAIL Pose Embeds",
    }