# pylint: disable=invalid-name, unused-import
"""Runtime stream APIs which are mainly for internal test use only.
For applications, please use PyTorch's stream management, of which DGL is aware.
"""
from __future__ import absolute_import

import ctypes
from .base import _LIB, check_call, _FFI_MODE
from .runtime_ctypes import DGLStreamHandle

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

def to_dgl_stream_handle(cuda_stream):
    """ Convert torch.cuda.Stream to DGL stream handle

    Parameters
    ----------
    cuda_stream : torch.cuda.Stream.

    Returns
    -------
    DGLStreamHandle
        DGLStreamHandle of the input ``cuda_stream``.
    """
    return ctypes.c_void_p(cuda_stream.cuda_stream)

def _dgl_get_stream(ctx):
    """Get the current CUDA stream of the given DGL context.

    Parameters
    ----------
    ctx : DGL context.

    Returns
    -------
    DGLStreamHandle
        DGLStreamHandle of the current CUDA stream.
    """
    current_cuda_stream = DGLStreamHandle()
    check_call(_LIB.DGLGetStream(
        ctx.device_type, ctx.device_id, ctypes.byref(current_cuda_stream)))
    return current_cuda_stream
