"""FikoYT DWPose Extract — comfyui_controlnet_aux DWPreprocessor sarmalayıcı."""
import os

_DWPOSE_NODE = None
try:
    from custom_nodes.comfyui_controlnet_aux.node_wrappers.dwpose import DWPose_Preprocessor as _dw  # type: ignore
    _DWPOSE_NODE = _dw
except Exception:
    try:
        # Bazı sürümlerde toplu import noktası
        from custom_nodes.comfyui_controlnet_aux import AIO_Preprocessor as _aio  # type: ignore
        _DWPOSE_NODE = _aio
    except Exception:
        pass

from ..core import ensure_model, get_models_dir


DWPOSE_MODELS = [
    {"repo": "yzd-v/DWPose",  "file": "dw-ll_ucoco_384.onnx", "subdir": "DWPose"},
    {"repo": "yzd-v/DWPose",  "file": "yolox_l.onnx",         "subdir": "DWPose"},
]


def _download_all(auto_download: bool) -> None:
    if not auto_download:
        return
    for m in DWPOSE_MODELS:
        target = os.path.join(get_models_dir(m["subdir"]), m["file"])
        if not os.path.isfile(target):
            ensure_model(m["repo"], m["file"], m["subdir"])


class FikoYTDWPoseExtract:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "include_body":  ("BOOLEAN", {"default": True}),
                "include_face":  ("BOOLEAN", {"default": True}),
                "include_hands": ("BOOLEAN", {"default": True}),
                "resolution":    ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
                "auto_download": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("pose_image",)
    FUNCTION = "extract"
    CATEGORY = "FikoYT/image"

    def extract(self, image, include_body, include_face, include_hands, resolution, auto_download):
        _download_all(auto_download)

        if _DWPOSE_NODE is None:
            raise RuntimeError(
                "comfyui_controlnet_aux DWPose node bulunamadı. "
                "https://github.com/Fannovel16/comfyui_controlnet_aux yükleyin."
            )

        try:
            instance = _DWPOSE_NODE()
            # DWPose_Preprocessor.estimate_pose(image, detect_hand, detect_body, detect_face, resolution, bbox_detector="yolox_l.onnx", pose_estimator="dw-ll_ucoco_384.onnx")
            if hasattr(instance, "estimate_pose"):
                out = instance.estimate_pose(
                    image=image,
                    detect_hand="enable" if include_hands else "disable",
                    detect_body="enable" if include_body else "disable",
                    detect_face="enable" if include_face else "disable",
                    resolution=resolution,
                    bbox_detector="yolox_l.onnx",
                    pose_estimator="dw-ll_ucoco_384.onnx",
                )
            else:
                raise AttributeError("DWPose estimate_pose API'si tanınmıyor.")
        except Exception as e:
            raise RuntimeError(
                f"DWPose çıkarımı başarısız: {e}\n"
                "comfyui_controlnet_aux > DWPose Estimator node'unu doğrudan deneyin."
            )

        if isinstance(out, dict):
            # AIO_Preprocessor bazen dict döner — IMAGE key'i al
            out = out.get("result", out)
        if isinstance(out, tuple):
            return (out[0],)
        return (out,)


NODE_CLASS_MAPPINGS = {"FikoYTDWPoseExtract": FikoYTDWPoseExtract}
NODE_DISPLAY_NAME_MAPPINGS = {"FikoYTDWPoseExtract": "FikoYT • DWPose Extract"}