# pylint: disable=invalid-name, unused-import
"""Function namespace."""
from __future__ import absolute_import

import ctypes
import sys

from .base import _FFI_MODE, _LIB, c_str, check_call, py_str, string_types

IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError

try:
    # pylint: disable=wrong-import-position
    if _FFI_MODE == "ctypes":
        raise ImportError()
    if sys.version_info >= (3, 0):
        from ._cy3.core import (
            _set_class_function,
            _set_class_module,
            convert_to_dgl_func,
            FunctionBase as _FunctionBase,
        )
    else:
        from ._cy2.core import (
            _set_class_function,
            _set_class_module,
            convert_to_dgl_func,
            FunctionBase as _FunctionBase,
        )
except IMPORT_EXCEPT:
    # pylint: disable=wrong-import-position
    from ._ctypes.function import (
        _set_class_function,
        _set_class_module,
        convert_to_dgl_func,
        FunctionBase as _FunctionBase,
    )

FunctionHandle = ctypes.c_void_p


class Function(_FunctionBase):
    """The PackedFunc object.

    Function plays an key role to bridge front and backend in DGL.
    Function provide a type-erased interface, you can call function with positional arguments.

    The compiled module returns Function.
    DGL backend also registers and exposes its API as Functions.
    For example, the developer function exposed in dgl.ir_pass are actually
    C++ functions that are registered as PackedFunc

    The following are list of common usage scenario of dgl.Function.

    - Automatic exposure of C++ API into python
    - To call PackedFunc from python side
    - To call python callbacks to inspect results in generated code
    - Bring python hook into C++ backend

    See Also
    --------
    dgl.register_func: How to register global function.
    dgl.get_global_func: How to get global function.
    """

    pass  # pylint: disable=unnecessary-pass


class ModuleBase(object):
    """Base class for module"""

    __slots__ = ["handle", "_entry", "entry_name"]

    def __init__(self, handle):
        self.handle = handle
        self._entry = None
        self.entry_name = "__dgl_main__"

    def __del__(self):
        check_call(_LIB.DGLModFree(self.handle))

    @property
    def entry_func(self):
        """Get the entry function

        Returns
        -------
        f : Function
            The entry function if exist
        """
        if self._entry:
            return self._entry
        self._entry = self.get_function(self.entry_name)
        return self._entry

    def get_function(self, name, query_imports=False):
        """Get function from the module.

        Parameters
        ----------
        name : str
            The name of the function

        query_imports : bool
            Whether also query modules imported by this module.

        Returns
        -------
        f : Function
            The result function.
        """
        ret_handle = FunctionHandle()
        check_call(
            _LIB.DGLModGetFunction(
                self.handle,
                c_str(name),
                ctypes.c_int(query_imports),
                ctypes.byref(ret_handle),
            )
        )
        if not ret_handle.value:
            raise AttributeError("Module has no function '%s'" % name)
        return Function(ret_handle, False)

    def import_module(self, module):
        """Add module to the import list of current one.

        Parameters
        ----------
        module : Module
            The other module.
        """
        check_call(_LIB.DGLModImport(self.handle, module.handle))

    def __getitem__(self, name):
        if not isinstance(name, string_types):
            raise ValueError("Can only take string as function name")
        return self.get_function(name)

    def __call__(self, *args):
        if self._entry:
            return self._entry(*args)
        f = self.entry_func
        return f(*args)


def register_func(func_name, f=None, override=False):
    """Register global function

    Parameters
    ----------
    func_name : str or function
        The function name

    f : function, optional
        The function to be registered.

    override: boolean optional
        Whether override existing entry.

    Returns
    -------
    fregister : function
        Register function if f is not specified.

    Examples
    --------
    The following code registers my_packed_func as global function.
    Note that we simply get it back from global function table to invoke
    it from python side. However, we can also invoke the same function
    from C++ backend, or in the compiled DGL code.

    .. code-block:: python

      targs = (10, 10.0, "hello")
      @dgl.register_func
      def my_packed_func(*args):
          assert(tuple(args) == targs)
          return 10
      # Get it out from global function table
      f = dgl.get_global_func("my_packed_func")
      assert isinstance(f, dgl.nd.Function)
      y = f(*targs)
      assert y == 10
    """
    if callable(func_name):
        f = func_name
        func_name = f.__name__

    if not isinstance(func_name, str):
        raise ValueError("expect string function name")

    ioverride = ctypes.c_int(override)

    def register(myf):
        """internal register function"""
        if not isinstance(myf, Function):
            myf = convert_to_dgl_func(myf)
        check_call(
            _LIB.DGLFuncRegisterGlobal(c_str(func_name), myf.handle, ioverride)
        )
        return myf

    if f:
        return register(f)
    return register


def get_global_func(name, allow_missing=False):
    """Get a global function by name

    Parameters
    ----------
    name : str
        The name of the global function

    allow_missing : bool
        Whether allow missing function or raise an error.

    Returns
    -------
    func : dgl.Function
        The function to be returned, None if function is missing.
    """
    handle = FunctionHandle()
    check_call(_LIB.DGLFuncGetGlobal(c_str(name), ctypes.byref(handle)))
    if handle.value:
        return Function(handle, False)
    else:
        if allow_missing:
            return None
        else:
            raise ValueError("Cannot find global function %s" % name)


def list_global_func_names():
    """Get list of global functions registered.

    Returns
    -------
    names : list
       List of global functions names.
    """
    plist = ctypes.POINTER(ctypes.c_char_p)()
    size = ctypes.c_uint()

    check_call(
        _LIB.DGLFuncListGlobalNames(ctypes.byref(size), ctypes.byref(plist))
    )
    fnames = []
    for i in range(size.value):
        fnames.append(py_str(plist[i]))
    return fnames


def extract_ext_funcs(finit):
    """
    Extract the extension PackedFuncs from a C module.

    Parameters
    ----------
    finit : ctypes function
        a ctypes that takes signature of DGLExtensionDeclarer

    Returns
    -------
    fdict : dict of str to Function
        The extracted functions
    """
    fdict = {}

    def _list(name, func):
        fdict[name] = func

    myf = convert_to_dgl_func(_list)
    ret = finit(myf.handle)
    _ = myf
    if ret != 0:
        raise RuntimeError("cannot initialize with %s" % finit)
    return fdict


def _get_api(f):
    flocal = f
    flocal.is_global = True
    return flocal


def _init_api(namespace, target_module_name=None):
    """Initialize api for a given module name

    namespace : str
       The namespace of the source registry

    target_module_name : str
       The target module name if different from namespace
    """
    target_module_name = target_module_name if target_module_name else namespace
    if namespace.startswith("dgl."):
        return _init_api_prefix(target_module_name, namespace[4:])
    else:
        return _init_api_prefix(target_module_name, namespace)


def _init_api_prefix(module_name, prefix):
    module = sys.modules[module_name]
    name_list = []

    for name in list_global_func_names():
        if name.startswith("_") and not name.startswith("_deprecate"):
            # internal APIs are ignored
            continue
        name_split = name.rsplit(".", 1)
        if name_split[0] != prefix:
            continue

        if len(name_split) == 1:
            print('Warning: invalid API name "%s".' % name)
            continue
        fname = name_split[1]
        target_module = module

        f = get_global_func(name)
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = "DGL PackedFunc %s. " % fname
        setattr(target_module, ff.__name__, ff)
        name_list.append(fname)

    return name_list


def _init_internal_api():
    for name in list_global_func_names():
        if not name.startswith("_") or name.startswith("_deprecate"):
            # normal APIs are ignored
            continue
        target_module = sys.modules["dgl._api_internal"]
        fname = name
        if fname.find(".") != -1:
            print('Warning: invalid API name "%s".' % fname)
            continue
        f = get_global_func(name)
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = "DGL PackedFunc %s. " % fname
        setattr(target_module, ff.__name__, ff)


_set_class_function(Function)
