# Copyright (c) OpenMMLab. All rights reserved.
import importlib
import os
import pkgutil
import warnings
from collections import namedtuple

import torch

if torch.__version__ != "parrots":

    def load_ext(name, funcs):
        ext = importlib.import_module("hordelib.nodes.comfy_controlnet_preprocessors.uniformer.mmcv." + name)
        for fun in funcs:
            assert hasattr(ext, fun), f"{fun} miss in module {name}"
        return ext

else:
    from parrots import extension
    from parrots.base import ParrotsException

    has_return_value_ops = [
        "nms",
        "softnms",
        "nms_match",
        "nms_rotated",
        "top_pool_forward",
        "top_pool_backward",
        "bottom_pool_forward",
        "bottom_pool_backward",
        "left_pool_forward",
        "left_pool_backward",
        "right_pool_forward",
        "right_pool_backward",
        "fused_bias_leakyrelu",
        "upfirdn2d",
        "ms_deform_attn_forward",
        "pixel_group",
        "contour_expand",
    ]

    def get_fake_func(name, e):
        def fake_func(*args, **kwargs):
            warnings.warn(f"{name} is not supported in parrots now")
            raise e

        return fake_func

    def load_ext(name, funcs):
        ExtModule = namedtuple("ExtModule", funcs)
        ext_list = []
        lib_root = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
        for fun in funcs:
            try:
                ext_fun = extension.load(fun, name, lib_dir=lib_root)
            except ParrotsException as e:
                if "No element registered" not in e.message:
                    warnings.warn(e.message)
                ext_fun = get_fake_func(fun, e)
                ext_list.append(ext_fun)
            else:
                if fun in has_return_value_ops:
                    ext_list.append(ext_fun.op)
                else:
                    ext_list.append(ext_fun.op_)
        return ExtModule(*ext_list)


def check_ops_exist():
    ext_loader = pkgutil.find_loader("mmcv._ext")
    return ext_loader is not None
