# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING

from cuda.bindings cimport cydriver
from cuda.core._device_resources cimport DeviceResources, SMResource, WorkqueueResource
from cuda.core._device_resources import SMResource, WorkqueueResource
from cuda.core._resource_handles cimport (
    ContextHandle,
    GreenCtxHandle,
    as_cu,
    create_context_handle_from_green_ctx,
    get_context_green_ctx,
    get_last_error,
    as_intptr,
    as_py,
)
from cuda.core._stream import Stream
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN

if TYPE_CHECKING:
    import cuda.bindings.driver  # no-cython-lint


__all__ = ['Context', 'ContextOptions']


DeviceResourcesType = Sequence[SMResource | WorkqueueResource]


cdef class Context:
    """CUDA context wrapper.

    Context objects represent CUDA contexts and cannot be instantiated directly.
    Use Device or Stream APIs to obtain context objects.
    """

    def __init__(self, *args, **kwargs) -> None:
        raise RuntimeError("Context objects cannot be instantiated directly. Please use Device or Stream APIs.")

    @staticmethod
    cdef Context _from_handle(type cls, ContextHandle h_context, int device_id):
        """Create Context from existing ContextHandle (cdef-only factory)."""
        cdef Context ctx = cls.__new__(cls)
        ctx._h_context = h_context
        ctx._device_id = device_id
        return ctx

    @staticmethod
    cdef Context _from_green_ctx(type cls, GreenCtxHandle h_green_ctx, int device_id):
        """Create Context from an owning green context handle."""
        cdef ContextHandle h_context = create_context_handle_from_green_ctx(h_green_ctx)
        if not h_context:
            HANDLE_RETURN(get_last_error())
            raise RuntimeError("Failed to create CUDA context view from green context")
        return Context._from_handle(cls, h_context, device_id)

    @property
    def handle(self) -> cuda.bindings.driver.CUcontext | None:
        """Return the underlying CUcontext handle."""
        if not self._h_context:
            return None
        if as_cu(self._h_context) == NULL:
            return None
        return as_py(self._h_context)

    @property
    def _handle(self) -> cuda.bindings.driver.CUcontext | None:
        return self.handle

    @property
    def is_green(self) -> bool:
        """True if this context was created from device resources."""
        if not self._h_context:
            return False
        return get_context_green_ctx(self._h_context).get() != NULL

    @property
    def resources(self) -> DeviceResources:
        """Query the hardware resources provisioned for this context.

        For green contexts, returns the resources this context was created
        with (SM partition, workqueue config). For primary contexts, returns
        the full device resources.

        Raises :class:`RuntimeError` if the context has been closed.
        """
        if not self._h_context:
            raise RuntimeError("Cannot query resources on a closed context")
        return DeviceResources._init_from_ctx(self._h_context, self._device_id)

    def create_stream(self, options: object = None) -> Stream:
        """Create a new stream bound to this green context.

        This method is only available on green contexts. For primary
        contexts, use :meth:`Device.create_stream` instead.

        Parameters
        ----------
        options : :obj:`~_stream.StreamOptions`, optional
            Customizable dataclass for stream creation options.

        Returns
        -------
        :obj:`~_stream.Stream`
            Newly created stream object.
        """
        if not self._h_context:
            raise RuntimeError("Cannot create a stream on a closed context")
        if not self.is_green:
            raise RuntimeError(
                "Context.create_stream() is only supported on green contexts. "
                "Use Device.create_stream() for primary contexts."
            )

        return Stream._init(options=options, device_id=self._device_id, ctx=self)

    cpdef close(self):
        """Release this context wrapper's underlying CUDA handles."""
        cdef cydriver.CUcontext current_ctx
        if self._h_context and as_cu(self._h_context) != NULL:
            with nogil:
                HANDLE_RETURN(cydriver.cuCtxGetCurrent(&current_ctx))
            if current_ctx == as_cu(self._h_context):
                raise RuntimeError(
                    "Cannot close a CUDA context while it is current. "
                    "Restore a previous context before closing this context."
                )
        self._h_context.reset()

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Context):
            return NotImplemented
        cdef Context _other = <Context>other
        return as_intptr(self._h_context) == as_intptr(_other._h_context)

    def __hash__(self) -> int:
        return hash(as_intptr(self._h_context))

    def __repr__(self) -> str:
        return f"<Context handle={as_intptr(self._h_context):#x} device={self._device_id}>"


@dataclass
cdef class ContextOptions:
    """Options for context creation.

    Attributes
    ----------
    resources : :obj:`~cuda.core.typing.DeviceResourcesType`
        Device resources used to create a green context.
    """
    resources: DeviceResourcesType
