"""
XB-ToolBox 人物分割节点 (ROCm/DirectML 优化)
============================================
完全独立，零外部节点依赖。使用 onnxruntime-directml 在 AMD 显卡上 GPU 运行。

配套节点:
  XB_HumanSegModelLoader  → 加载 ONNX 模型 (文件选择器)
  XB_HumanSegmentation    → 人物分割 (蒙版 + 抠图)

模型下载: https://huggingface.co/briaai/RMBG-1.4
放入: ComfyUI/models/rembg/u2net_human_seg.onnx
"""

import os
import numpy as np
import torch
import folder_paths


# 注册 rembg 模型文件夹
if "rembg" not in folder_paths.folder_names_and_paths:
    folder_paths.folder_names_and_paths["rembg"] = (
        [os.path.join(folder_paths.models_dir, "rembg")],
        {".onnx"}
    )


def _get_ort_session(model_path: str):
    """获取 ONNX Runtime 会话，优先 DirectML GPU，回退 CPU"""
    import onnxruntime as ort

    providers = ort.get_available_providers()
    if "DmlExecutionProvider" in providers:
        print(f"[XB-Seg] ✅ DirectML GPU 加速已启用")
        return ort.InferenceSession(model_path, providers=["DmlExecutionProvider"])
    elif "ROCMExecutionProvider" in providers:
        print(f"[XB-Seg] ✅ ROCm GPU 加速已启用")
        return ort.InferenceSession(model_path, providers=["ROCMExecutionProvider"])
    elif "CUDAExecutionProvider" in providers:
        print(f"[XB-Seg] ✅ CUDA GPU 加速已启用")
        return ort.InferenceSession(model_path, providers=["CUDAExecutionProvider"])
    else:
        print(f"[XB-Seg] ⚠️ 未检测到 GPU 加速器，使用 CPU (速度较慢)")
        return ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])


def _preprocess(image_bchw: torch.Tensor) -> np.ndarray:
    """[B, 3, 320, 320] → numpy, ImageNet 标准化"""
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 3, 1, 1)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 3, 1, 1)
    x = image_bchw.cpu().numpy().astype(np.float32)
    return (x - mean) / std


def _postprocess(output: np.ndarray, orig_h: int, orig_w: int) -> torch.Tensor:
    """ONNX 输出 [1,1,320,320] → resize 回原始尺寸 → [1, H, W] mask"""
    import cv2
    mask = output[0, 0]  # [320, 320]
    mask = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
    mask = np.clip(mask, 0, 1)
    return torch.from_numpy(mask).unsqueeze(0)


# ============================================================
# 节点 1: 模型加载器
# ============================================================
class XB_HumanSegModelLoader:
    """加载 u2net ONNX 模型，输出给分割节点使用"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": (folder_paths.get_filename_list("rembg"),),
            }
        }

    RETURN_TYPES = ("XB_HUMANSEG_MODEL",)
    RETURN_NAMES = ("分割模型",)
    FUNCTION = "load"
    CATEGORY = "XB_ToolBox/Segment"

    def load(self, model: str):
        model_path = folder_paths.get_full_path_or_raise("rembg", model)
        session = _get_ort_session(model_path)
        print(f"[XB-Seg] 已加载: {model}")
        return (session,)


# ============================================================
# 节点 2: 人物分割
# ============================================================
class XB_HumanSegmentation:
    """输入图像 + 分割模型，输出人物蒙版和抠图"""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "seg_model": ("XB_HUMANSEG_MODEL",),
                "image": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE")
    RETURN_NAMES = ("人物蒙版", "人物抠图")
    FUNCTION = "segment"
    CATEGORY = "XB_ToolBox/Segment"

    def segment(self, seg_model, image: torch.Tensor):
        """
        seg_model: onnxruntime.InferenceSession
        image: [B, H, W, C] float32 [0, 1]
        """
        session = seg_model
        B, orig_h, orig_w, C = image.shape

        # --- 1. 预处理: resize 到 320x320，BCHW ---
        img_bchw = image.permute(0, 3, 1, 2)  # [B, C, H, W]
        img_bchw = torch.nn.functional.interpolate(
            img_bchw, size=(320, 320), mode="bilinear", align_corners=False
        )
        inp = _preprocess(img_bchw)  # numpy [B, 3, 320, 320]

        # --- 2. ONNX 推理 ---
        ort_inputs = {session.get_inputs()[0].name: inp}
        ort_outputs = session.run(None, ort_inputs)
        mask_np = ort_outputs[0]  # [B, 1, 320, 320]

        # --- 3. 后处理 ---
        masks = []
        cutouts = []
        for b_idx in range(B):
            mask = _postprocess(mask_np[b_idx:b_idx + 1], orig_h, orig_w)  # [1, H, W]
            masks.append(mask)
            fg = image[b_idx] * mask.permute(1, 2, 0)  # [H, W, C]
            cutouts.append(fg)

        mask_stack = torch.cat(masks, dim=0)      # [B, H, W]
        cutout_stack = torch.stack(cutouts, dim=0)  # [B, H, W, C]

        return (mask_stack, cutout_stack)


__all__ = ["XB_HumanSegModelLoader", "XB_HumanSegmentation"]