# -----------------------------------------------------------------
# 这是一个ComfyUI的自定义节点
# 功能：为图像添加一个双行、可自定义字体大小的水印。
# 作者：(根据您的函数创建)
# -----------------------------------------------------------------

import torch
import numpy as np
from PIL import Image, ImageDraw, ImageFont

# -----------------------------------------------------------------
# 核心水印添加函数 (由用户提供)
# -----------------------------------------------------------------
def add_watermark(image, font_path,
                  watermark_text_line1="立香", font_size_line1=40,
                  watermark_text_line2="美羽", font_size_line2=30):
    """
    Adds a two-line watermark to an image with individually adjustable font sizes.
    """
    if not isinstance(image, Image.Image):
        raise ValueError("Input must be a PIL Image object")

    if image.mode != 'RGBA':
        image = image.convert('RGBA')

    width, height = image.size
    draw = ImageDraw.Draw(image)

    try:
        font1 = ImageFont.truetype(font_path, font_size_line1)
    except IOError:
        print(f"警告：找不到字体文件 {font_path}。将使用默认字体。")
        font1 = ImageFont.load_default()

    try:
        font2 = ImageFont.truetype(font_path, font_size_line2)
    except IOError:
        font2 = ImageFont.load_default()

    bbox1 = draw.textbbox((0, 0), watermark_text_line1, font=font1)
    text_width_line1 = bbox1[2] - bbox1[0]
    text_height_line1 = bbox1[3] - bbox1[1]

    bbox2 = draw.textbbox((0, 0), watermark_text_line2, font=font2)
    text_width_line2 = bbox2[2] - bbox2[0]
    text_height_line2 = bbox2[3] - bbox2[1]

    margin = 20
    line_spacing = 10

    y_line2 = height - text_height_line2 - margin
    x_line2 = width - text_width_line2 - margin

    y_line1 = y_line2 - text_height_line1 - line_spacing
    x_line1 = width - text_width_line1 - margin

    # --- 这里是唯一的修改 ---
    # 将水印颜色设置为纯白色 (R, G, B, Alpha)
    # Alpha = 255 代表完全不透明
    white_color = (255, 255, 255, 255) 

    draw.text((x_line1, y_line1), watermark_text_line1, font=font1, fill=white_color)
    draw.text((x_line2, y_line2), watermark_text_line2, font=font2, fill=white_color)

    return image

# -----------------------------------------------------------------
# ComfyUI 节点类
# -----------------------------------------------------------------
class WatermarkNode:
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义节点的输入"""
        return {
            "required": {
                "image": ("IMAGE",),
                "font_path": ("STRING", {
                    "multiline": False,
                    # 为Windows用户提供一个常见的中文字体路径作为默认值
                    # Linux/Mac用户需要修改此路径
                    "default": "D:\\ComfyUI_windows_portable\\ComfyUI\\Fonts\\Iansui-Regular.ttf" 
                }),
                "text_line1": ("STRING", {"multiline": False, "default": "妃妃"}),
                "font_size_line1": ("INT", {"default": 40, "min": 1, "max": 1024, "step": 1}),
                "text_line2": ("STRING", {"multiline": False, "default": "aiFeiFei"}),
                "font_size_line2": ("INT", {"default": 30, "min": 1, "max": 1024, "step": 1}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_watermark"
    CATEGORY = "Image/Post-Processing"

    def apply_watermark(self, image, font_path, text_line1, font_size_line1, text_line2, font_size_line2):
        """节点的核心执行逻辑"""
        
        # ComfyUI 的图像输入是 PyTorch 张量，我们需要将其转换为 PIL 图像
        # 形状: [批次数, 高度, 宽度, 通道数], 值范围: 0.0-1.0
        
        watermarked_images = []
        for i in range(image.shape[0]):
            # 1. 从批次中取出一个张量
            img_tensor = image[i]
            
            # 2. 将张量转换为 NumPy 数组并调整数值范围 (0-1 -> 0-255)
            img_np = np.clip(255. * img_tensor.cpu().numpy(), 0, 255).astype(np.uint8)
            
            # 3. 从 NumPy 数组创建 PIL 图像
            pil_image = Image.fromarray(img_np)

            # 4. 调用核心函数添加水印
            pil_image_watermarked = add_watermark(
                pil_image,
                font_path,
                text_line1,
                font_size_line1,
                text_line2,
                font_size_line2
            )
            
            # 5. 将处理后的 PIL 图像转换回 NumPy 数组
            img_np_watermarked = np.array(pil_image_watermarked).astype(np.float32)
            
            # 6. 将数值范围调回 0-1 并转换回 PyTorch 张量
            img_tensor_watermarked = torch.from_numpy(img_np_watermarked / 255.0)
            
            watermarked_images.append(img_tensor_watermarked)

        # 将处理后的图像列表堆叠成一个批次张量
        final_tensor = torch.stack(watermarked_images)

        return (final_tensor,)

# -----------------------------------------------------------------
# ComfyUI 必须的映射字典
# -----------------------------------------------------------------
NODE_CLASS_MAPPINGS = {
    "WatermarkNode": WatermarkNode
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "WatermarkNode": "图像水印 (Watermark)"
}