import math
import torch
import comfy.utils

class WanNativeResize_Dolphin:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "target_model": (["480P (832x480)", "720P (1280x720)", "1080P (1920x1080)"], {"default": "720P (1280x720)"}),
                "upscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], {"default": "lanczos"}),
                "alignment": ("INT", {"default": 16, "min": 8, "max": 64, "step": 8}),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT")
    RETURN_NAMES = ("IMAGE", "width", "height")
    FUNCTION = "resize_for_wan"
    CATEGORY = "Dolphin Node/Wan Video" 

    def resize_for_wan(self, image, target_model, upscale_method, alignment):
        # 1. ComfyUI 이미지 텐서의 형태를 가져옵니다: (Batch/Frames, Height, Width, Channels)
        b, h, w, c = image.shape
        
        # 2. 타겟 모델에 따른 총 픽셀 면적(Area) 세팅
        if "480P" in target_model:
            target_area = 832 * 480
        elif "720P" in target_model:
            target_area = 1280 * 720
        else:
            target_area = 1920 * 1080
            
        # 3. Byungjoo님의 수학 로직 (비율 유지 계산)
        aspect_ratio = w / h
        new_h = math.sqrt(target_area / aspect_ratio)
        new_w = new_h * aspect_ratio
        
        # 4. [핵심] VAE가 좋아하는 배수(Alignment, 보통 16)로 반올림 처리
        new_w = int(round(new_w / alignment) * alignment)
        new_h = int(round(new_h / alignment) * alignment)

        # 5. [최적화] 이미 목표 해상도와 일치한다면, 무거운 연산 없이 원본 통과 (패스스루)
        if w == new_w and h == new_h:
            return (image, new_w, new_h)

        # 6. [핵심] 텐서 차원 변경: ComfyUI (B, H, W, C) -> 리사이즈 엔진용 (B, C, H, W)
        image = image.movedim(-1, 1)

        # 7. ComfyUI 내장 엔진을 사용해 초고화질 리사이즈 (Lanczos 지원)
        resized_image = comfy.utils.common_upscale(image, new_w, new_h, upscale_method, "disabled")

        # 8. 텐서 차원 원상 복구: (B, C, H, W) -> ComfyUI (B, H, W, C)
        resized_image = resized_image.movedim(1, -1)

        return (resized_image, new_w, new_h)