"""CuPy External Memory Manager plugins for Numba CUDA contexts.

This module adapts CuPy synchronous and asynchronous memory pools for Numba's
External Memory Manager interface. It also exposes helpers that align CuPy's
current stream selection with streams generated by :mod:`numba.cuda` so that
stream-ordered allocations behave consistently across both libraries.
"""

from contextlib import contextmanager
import logging
from types import TracebackType
from typing import Any, Callable, Iterator, Optional

from numba import cuda
import ctypes

from cubie.cuda_simsafe import (
    GetIpcHandleMixin,
    HostOnlyCUDAMemoryManager,
    MemoryPointer,
    MemoryInfo,
    Stream
)


logger = logging.getLogger(__name__)


def _numba_stream_ptr(
    nb_stream: Optional[Stream],
) -> Optional[int]:
    """
    Extract a ``CUstream`` pointer from a Numba stream wrapper.

    Parameters
    ----------
    nb_stream
        Numba CUDA stream whose ``CUstream`` pointer should be extracted. When
        ``None``, pointer extraction is skipped.

    Returns
    -------
    int or None
        Pointer value compatible with CuPy external streams, or ``None`` when
        extraction fails.

    Notes
    -----
    The function checks common attribute layouts across supported Numba
    versions to maintain compatibility.
    """
    if nb_stream is None:
        return None
    h = getattr(nb_stream, "handle", None)
    if h is None:
        return None
    # ctypes.c_void_p or int-like
    if isinstance(h, ctypes.c_void_p):
        return int(h.value) if h.value is not None else None
    try:
        return int(getattr(h, "value", h))
    except Exception:
        return None


class current_cupy_stream:
    """Context manager that forwards a Numba stream into CuPy APIs.

    Parameters
    ----------
    nb_stream
        Numba CUDA stream to expose to CuPy.

    Attributes
    ----------
    nb_stream
        The Numba stream being forwarded.
    cupy_ext_stream
        CuPy external stream wrapper around the Numba stream when CuPy manages
        the active allocator.

    Notes
    -----
    This context manager only affects CuPy-backed memory managers. When the
    active manager is not CuPy-based, the context operates as a no-op.
    """

    def __init__(self, nb_stream: Stream) -> None:
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        self.nb_stream = nb_stream

        self.cupy_ext_stream = None
        try:
            self._mgr_is_cupy = cuda.current_context().memory_manager.is_cupy
        except AttributeError:  # Numba allocators have no such attribute
            self._mgr_is_cupy = False

    def __enter__(self) -> "current_cupy_stream":
        """
        Enter the context and set up a CuPy external stream if applicable.

        Returns
        -------
        current_cupy_stream
            The active context manager instance.
        """
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        if self._mgr_is_cupy:
            ptr = _numba_stream_ptr(self.nb_stream)
            if ptr:
                self.cupy_ext_stream = cp.cuda.ExternalStream(ptr)
                self.cupy_ext_stream.__enter__()
            return self
        else:
            return self

    def __exit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        """
        Exit the context and clean up the CuPy external stream.

        Parameters
        ----------
        exc_type
            Exception type if an exception occurred.
        exc
            Exception instance if an exception occurred.
        tb
            Traceback object if an exception occurred.

        Returns
        -------
        None
            ``None``.
        """
        if self._mgr_is_cupy:
            if self.cupy_ext_stream is not None:
                self.cupy_ext_stream.__exit__(exc_type, exc, tb)
                self.cupy_ext_stream = None


class CuPyNumbaManager(GetIpcHandleMixin, HostOnlyCUDAMemoryManager):
    """
    Base Numba EMM plugin for using CuPy memory pools to allocate.

    Parameters
    ----------
    context
        CUDA context for memory management.

    Attributes
    ----------
    is_cupy
        Flag indicating this is a CuPy-based memory manager.

    Notes
    -----
    Drawn from the tutorial example at:
    https://github.com/numba/nvidia-cuda-tutorial/blob/main/session-5/examples/cupy_emm_plugin.py

    Extended to handle passing numba-generated streams as CuPy external
    streams, such that the allocations are stream-ordered when using the async
    allocator.
    """

    def __init__(self, context) -> None:
        try:
            import cupy as cp
        except ImportError:
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        super().__init__(context=context)
        # We keep a record of all allocations, and remove each allocation
        # record in the finalizer, which results in it being returned back to
        # the CuPy memory pool.
        self._allocations = {}
        # The CuPy memory pool.
        self._mp = None
        self._ctx = context
        self.is_cupy = True

        # These provide a way for tests to check who's allocating what.
        self._testing = False
        self._testout = None

    def memalloc(self, nbytes: int) -> MemoryPointer:
        """
        Allocate memory from the CuPy pool.

        Parameters
        ----------
        nbytes
            Number of bytes to allocate.

        Returns
        -------
        MemoryPointer
            Numba memory pointer wrapping the CuPy allocation.
        """
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        # Allocate from the CuPy pool and wrap the result in a MemoryPointer as
        # required by Numba.
        cp_mp = self._mp.malloc(nbytes)
        logger.debug("Allocated %d bytes at %x" % (nbytes, cp_mp.ptr))
        logger.debug("on stream %s" % (cp.cuda.get_current_stream()))
        self._allocations[cp_mp.ptr] = cp_mp
        return MemoryPointer(
            cuda.current_context(),
            ctypes.c_void_p(int(cp_mp.ptr)),
            nbytes,
            finalizer=self._make_finalizer(cp_mp, nbytes),
        )

    def _make_finalizer(self, cp_mp: Any, nbytes: int) -> Callable[[], None]:
        """
        Create a finalizer function for memory cleanup.

        Parameters
        ----------
        cp_mp
            CuPy memory pool allocation to be cleaned up.
        nbytes
            Number of bytes in the allocation.

        Returns
        -------
        Callable
            Finalizer function that removes the allocation reference.
        """
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        allocations = self._allocations
        ptr = cp_mp.ptr

        def finalizer():
            logger.debug("Freeing %d bytes at %x" % (nbytes, ptr))
            logger.debug("on stream %s" % (cp.cuda.get_current_stream()))
            # Removing the last reference to the allocation causes it to be
            # garbage-collected and returned to the pool.
            allocations.pop(ptr)

        return finalizer

    def get_memory_info(self) -> MemoryInfo:
        """
        Get memory information from the CuPy memory pool.

        Returns
        -------
        MemoryInfo
            Object containing free and total memory in bytes from the pool.

        Notes
        -----
        Returns information from the CuPy memory pool, not the whole device.
        """
        return MemoryInfo(
            free=self._mp.free_bytes(), total=self._mp.total_bytes()
        )

    def initialize(self) -> None:
        """
        Initialize the memory manager.

        Returns
        -------
        None
            ``None``.
        """
        super().initialize()

    def reset(self, stream: Optional[Any] = None) -> None:
        """
        Free all blocks with optional stream for async operations.

        Parameters
        ----------
        stream
            CuPy stream used for asynchronous deallocation. When omitted, the
            pool frees blocks synchronously.

        Returns
        -------
        None
            ``None``.

        Notes
        -----
        This is called without a stream argument when the context is reset. To
        run the operation in one stream, call this function by itself using
        :code:`cuda.current_context().memory_manager.reset(stream)`.
        """
        super().reset()
        if self._mp:
            self._mp.free_all_blocks(stream=stream)

    @contextmanager
    def defer_cleanup(self) -> Iterator[None]:
        """
        Context manager for deferring memory cleanup operations.

        Returns
        -------
        Iterator
            Generator yielded by :func:`contextlib.contextmanager` for use in a
            ``with`` block that yields ``None`` to the caller.

        Notes
        -----
        This does not actually defer returning memory back to the pool, but
        returning memory to the pool will not interrupt async operations like
        an actual :func:`numba.cuda.cudadrv.driver.device_free` call would.
        """
        with super().defer_cleanup():
            yield

    @property
    def interface_version(self):
        """EMM interface version number."""
        return 1


class CuPyAsyncNumbaManager(CuPyNumbaManager):
    """
    Numba EMM plugin using CuPy MemoryAsyncPool for allocation and freeing.

    Parameters
    ----------
    context
        CUDA context for memory management.

    Notes
    -----
    Uses CuPy's asynchronous memory pool which provides stream-ordered
    memory operations.
    """

    def __init__(self, context) -> None:
        super().__init__(context=context)

    def initialize(self) -> None:
        """
        Initialize the async memory pool.

        Returns
        -------
        None
            ``None``.
        """
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        super().initialize()
        self._mp = cp.cuda.MemoryAsyncPool()

    def memalloc(self, nbytes: int) -> MemoryPointer:
        """
        Allocate memory from the async CuPy pool.

        Parameters
        ----------
        nbytes
            Number of bytes to allocate.

        Returns
        -------
        MemoryPointer
            Numba memory pointer wrapping the CuPy allocation.
        """
        if self._testing:
            self._testout = "async"
        return super().memalloc(nbytes)


class CuPySyncNumbaManager(CuPyNumbaManager):
    """
    Numba EMM plugin using CuPy MemoryPool for allocation and freeing.

    Parameters
    ----------
    context
        CUDA context for memory management.

    Notes
    -----
    Uses CuPy's synchronous memory pool which provides standard
    memory operations.
    """

    def __init__(self, context):
        super().__init__(context=context)

    def initialize(self) -> None:
        """
        Initialize the sync memory pool.

        Returns
        -------
        None
            ``None``.
        """
        try:
            import cupy as cp
        except ImportError: # pragma: no cover
            raise ImportError(
                "To use Cupy memory managers, you must install cupy: pip "
                "install cupy-cuda12x (assuming CUDA toolkit 12.x installed)]"
            )
        super().initialize()
        # Get the default memory pool for this context.
        self._mp = cp.get_default_memory_pool()

    def memalloc(self, nbytes: int) -> MemoryPointer:
        """
        Allocate memory from the sync CuPy pool.

        Parameters
        ----------
        nbytes
            Number of bytes to allocate.

        Returns
        -------
        MemoryPointer
            Numba memory pointer wrapping the CuPy allocation.
        """
        if self._testing:
            self._testout = "sync"
        return super().memalloc(nbytes)
