import json

from .test_single_image import test_node


class FaceDetect:
    """
    人脸遮挡检测
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "main_seed": (
                    "INT:seed",
                    {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
                ),
                "model": (["convnext_tiny", "convnext_base"],),
                "length": ("INT", {"default": 10, "min": 3, "max": 60, "step": 1}),
                "threshold": (
                    "FLOAT",
                    {"default": 94, "min": 55, "max": 99, "step": 0.1},
                ),
            },
        }

    RETURN_TYPES = (
        "IMAGE",
        "IMAGE",
        "STRING",
        "STRING",
        "STRING",
        "STRING",
        "STRING",
        "INT",
        "INT",
    )
    RETURN_NAMES = (
        "图像",
        "选中人脸",
        "分类",
        "概率",
        "采用帧序号",
        "全部帧序列",
        "剪辑配置",
        "起始帧序号",
        "帧数量",
    )

    FUNCTION = "predict"

    CATEGORY = "不忘科技-自定义节点🚩/面部"

    def predict(self, image, main_seed, model, length, threshold):
        image, image_selected, cls, prob, nums, period = test_node(
            image, length=length, thres=threshold, model_name=model
        )
        print("全部帧序列", period)
        if len(period) > 0:
            start, end = period[main_seed % len(period)]
            config = {"start": start, "end": end}
        else:
            raise RuntimeError("未找到符合要求的视频片段")
        return (
            image,
            image_selected,
            cls,
            prob,
            nums,
            str(period),
            json.dumps(config),
            start,
            end - start + 1,
        )