"""直接调 ComfyUI 重绘 — 替代 local_test HTTP 服务。

将 local_test/app.py 的核心逻辑（遮罩处理 + ComfyUI 调用）提取为 Python 函数，
不再需要独立 Flask 服务。使用 0716add-hair-api.json 工作流（steps=4）。
"""
from __future__ import annotations

import io
import logging
import os

import numpy as np
from PIL import Image, ImageFilter

from . import comfyui

logger = logging.getLogger("hair.worker")

_DEFAULT_PROMPT = "填充遮罩区域的头发"
_REPO = os.path.dirname(os.path.dirname(__file__))
_REPAINT_WORKFLOW = os.path.join(_REPO, "0716add-hair-api.json")


def _process_mask_to_rgba(image_bytes: bytes, mask_bytes: bytes) -> bytes:
    """将分开的 image + mask 处理为 ComfyUI 用的 RGBA PNG bytes。

    复制 local_test/app.py 的遮罩处理逻辑：
    1. 加载 image 为 RGB
    2. 加载 mask 为 RGBA，取所有通道 max 值（支持红/白/alpha 遮罩）
    3. resize mask 到与 image 一致
    4. 高斯模糊(radius=4) 柔化边缘
    5. alpha = 255 - mask（绘制区=255 → alpha=0 → 重绘区）
    6. 合成 RGBA PNG
    """
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    mask_img = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
    mask_arr = np.array(mask_img)
    mask_data = np.max(mask_arr, axis=2)  # (H, W) uint8

    mask_data_img = Image.fromarray(mask_data, mode="L")
    if mask_data_img.size != image.size:
        mask_data_img = mask_data_img.resize(image.size, Image.LANCZOS)
    mask_data_img = mask_data_img.filter(ImageFilter.GaussianBlur(radius=4))

    # ComfyUI LoadImage: mask = 1.0 - (alpha/255)
    # alpha=0 -> mask=1.0 (inpaint), alpha=255 -> mask=0.0 (keep)
    comfyui_alpha = Image.eval(mask_data_img, lambda x: 255 - x)

    r, g, b = image.split()
    rgba = Image.merge("RGBA", (r, g, b, comfyui_alpha))

    buf = io.BytesIO()
    rgba.save(buf, format="PNG")
    return buf.getvalue()


def run_redraw(image_bytes: bytes, mask_bytes: bytes,
               prompt: str | None = None, timeout: float = 300.0,
               front: bool = False, unet_name: str | None = None) -> bytes:
    """直接调 ComfyUI 重绘 — 替代 local_test /api/generate。

    Args:
        image_bytes: 人物图片字节（JPG/PNG）
        mask_bytes: 遮罩图片字节（支持红/白/alpha 遮罩格式）
        prompt: 提示词，None 用默认 "填充遮罩区域的头发"
        timeout: ComfyUI 超时秒数
        front: True 时任务插到 ComfyUI 队列最前（接口2 时延敏感路径用）
        unet_name: 非 None 时切换 Flux 模型（如 flux-2-klein-9b-Q5_K_M.gguf），None 用工作流默认

    Returns:
        重绘后的 PNG 图片字节

    Raises:
        RuntimeError: ComfyUI 执行失败
        TimeoutError: ComfyUI 超时
    """
    rgba_png = _process_mask_to_rgba(image_bytes, mask_bytes)
    return comfyui.run(rgba_png, timeout=timeout, prompt=prompt,
                       workflow_path=_REPAINT_WORKFLOW, front=front,
                       unet_name=unet_name)