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

import functools
from functools import partial
import multiprocessing
import platform
import warnings
from collections import namedtuple
from collections.abc import Sequence
from contextlib import ExitStack
from typing import Any, Callable

from cuda.bindings import driver as driver, nvrtc as nvrtc, runtime as runtime

# Module-level annotations that reference `driver`, `nvrtc`, and `runtime` so
# that stubgen-pyx keeps these imports in the generated `.pyi` (it would
# otherwise trim them as unused). These names are not assigned, so they only
# affect __annotations__ and have no runtime cost.
_keep_driver_in_stub: 'driver.CUresult'
_keep_nvrtc_in_stub: 'nvrtc.nvrtcResult'
_keep_runtime_in_stub: 'runtime.cudaError_t'

from cuda.bindings.nvvm import nvvmError
from cuda.bindings.nvjitlink import nvJitLinkError

from cpython.buffer cimport PyObject_GetBuffer, PyBuffer_Release, Py_buffer, PyBUF_SIMPLE

from cuda.bindings cimport cynvrtc, cynvvm, cynvjitlink

from cuda.core._utils.driver_cu_result_explanations import DRIVER_CU_RESULT_EXPLANATIONS
from cuda.core._utils.runtime_cuda_error_explanations import RUNTIME_CUDA_ERROR_EXPLANATIONS


class CUDAError(Exception):
    pass


class NVRTCError(CUDAError):
    pass



ComputeCapability = namedtuple("ComputeCapability", ("major", "minor"))


def cast_to_3_tuple(label: str, cfg: int | tuple[int, ...]) -> tuple[int, int, int]:
    cfg_orig = cfg
    if isinstance(cfg, int):
        cfg = (cfg,)
    else:
        common = "must be an int, or a tuple with up to 3 ints"
        if not isinstance(cfg, tuple):
            raise ValueError(f"{label} {common} (got {type(cfg)})")
        if len(cfg) > 3:
            raise ValueError(f"{label} {common} (got tuple with length {len(cfg)})")
        if any(not isinstance(val, int) for val in cfg):
            raise ValueError(f"{label} {common} (got {cfg})")
    if any(val < 1 for val in cfg):
        plural_s = "" if len(cfg) == 1 else "s"
        raise ValueError(f"{label} value{plural_s} must be >= 1 (got {cfg_orig})")
    return cfg + (1,) * (3 - len(cfg))


cdef int HANDLE_RETURN(cydriver.CUresult err) except?-1 nogil:
    if err != cydriver.CUresult.CUDA_SUCCESS:
        return _check_driver_error(err)
    return 0


cdef int _get_current_device_id() except? -1:
    """Return the current thread's bound CUdevice ordinal."""
    cdef cydriver.CUdevice dev
    with nogil:
        HANDLE_RETURN(cydriver.cuCtxGetDevice(&dev))
    return <int>dev


cdef int HANDLE_RETURN_NVRTC(cynvrtc.nvrtcProgram prog, cynvrtc.nvrtcResult err) except?-1 nogil:
    """Handle NVRTC result codes, raising NVRTCError with program log on failure."""
    if err == cynvrtc.nvrtcResult.NVRTC_SUCCESS:
        return 0
    with gil:
        _raise_nvrtc_error(prog, err)


cdef int _raise_nvrtc_error(cynvrtc.nvrtcProgram prog, cynvrtc.nvrtcResult err) except -1:
    """Build error message with program log and raise NVRTCError."""
    cdef const char* err_str = cynvrtc.nvrtcGetErrorString(err)
    cdef size_t logsize = 0
    if prog != NULL:
        cynvrtc.nvrtcGetProgramLogSize(prog, &logsize)
    cdef bytes log_bytes
    cdef str log_str = ""
    if logsize > 1 and prog != NULL:
        log_bytes = b" " * logsize
        if cynvrtc.nvrtcGetProgramLog(prog, <char*>log_bytes) == cynvrtc.nvrtcResult.NVRTC_SUCCESS:
            log_str = log_bytes.decode("utf-8", errors="backslashreplace")
    err_msg = f"{err}: {err_str.decode()}" if err_str != NULL else f"NVRTC error {err}"
    if log_str:
        err_msg += f", compilation log:\n\n{log_str}"
    raise NVRTCError(err_msg)


cdef int HANDLE_RETURN_NVVM(cynvvm.nvvmProgram prog, cynvvm.nvvmResult err) except?-1 nogil:
    """Handle NVVM result codes, raising nvvmError with program log on failure."""
    if err == cynvvm.nvvmResult.NVVM_SUCCESS:
        return 0
    with gil:
        _raise_nvvm_error(prog, err)


cdef int _raise_nvvm_error(cynvvm.nvvmProgram prog, cynvvm.nvvmResult err) except -1:
    """Raise nvvmError annotated with the program log."""
    cdef size_t logsize = 0
    if prog != NULL:
        cynvvm.nvvmGetProgramLogSize(prog, &logsize)
    cdef bytes log_bytes
    cdef str log_str = ""
    if logsize > 1 and prog != NULL:
        log_bytes = b" " * logsize
        if cynvvm.nvvmGetProgramLog(prog, <char*>log_bytes) == cynvvm.nvvmResult.NVVM_SUCCESS:
            log_str = log_bytes.decode("utf-8", errors="backslashreplace")
    cdef object exc = nvvmError(err)
    if log_str:
        exc.args = (exc.args[0] + f"\nNVVM program log: {log_str}", *exc.args[1:])
    raise exc


cdef int HANDLE_RETURN_NVJITLINK(
        cynvjitlink.nvJitLinkHandle handle, cynvjitlink.nvJitLinkResult err) except?-1 nogil:
    """Handle nvJitLink result codes, raising nvJitLinkError with error log on failure."""
    if err == cynvjitlink.nvJitLinkResult.NVJITLINK_SUCCESS:
        return 0
    with gil:
        _raise_nvjitlink_error(handle, err)


cdef int _raise_nvjitlink_error(
        cynvjitlink.nvJitLinkHandle handle, cynvjitlink.nvJitLinkResult err) except -1:
    """Raise nvJitLinkError annotated with the error log."""
    cdef size_t logsize = 0
    if handle != NULL:
        cynvjitlink.nvJitLinkGetErrorLogSize(handle, &logsize)
    cdef bytes log_bytes
    cdef str log_str = ""
    if logsize > 1 and handle != NULL:
        log_bytes = b" " * logsize
        if cynvjitlink.nvJitLinkGetErrorLog(handle, <char*>log_bytes) == \
                cynvjitlink.nvJitLinkResult.NVJITLINK_SUCCESS:
            log_str = log_bytes.decode("utf-8", errors="backslashreplace")
    cdef object exc = nvJitLinkError(err)
    if log_str:
        exc.args = (exc.args[0] + f"\nnvJitLink error log: {log_str}", *exc.args[1:])
    raise exc


cdef object _RUNTIME_SUCCESS = runtime.cudaError_t.cudaSuccess
cdef object _NVRTC_SUCCESS = nvrtc.nvrtcResult.NVRTC_SUCCESS


cpdef inline int _check_driver_error(cydriver.CUresult error) except?-1 nogil:
    if error == cydriver.CUresult.CUDA_SUCCESS:
        return 0
    cdef const char* name
    name_err = cydriver.cuGetErrorName(error, &name)
    if name_err != cydriver.CUresult.CUDA_SUCCESS:
        raise CUDAError(f"UNEXPECTED ERROR CODE: {error}")
    with gil:
        # TODO: consider lower this to Cython
        expl = DRIVER_CU_RESULT_EXPLANATIONS.get(int(error))
        if expl is not None:
            raise CUDAError(f"{name.decode()}: {expl}")
    cdef const char* desc
    desc_err = cydriver.cuGetErrorString(error, &desc)
    if desc_err != cydriver.CUresult.CUDA_SUCCESS:
        raise CUDAError(f"{name.decode()}")
    raise CUDAError(f"{name.decode()}: {desc.decode()}")


cpdef inline int _check_runtime_error(error) except?-1:
    if error == _RUNTIME_SUCCESS:
        return 0
    # `_check_error()` reaches this path only for `runtime.cudaError_t` values.
    # Use the enum name directly because Windows hybrid cudart can lag that table.
    name = error.name
    expl = RUNTIME_CUDA_ERROR_EXPLANATIONS.get(int(error))
    if expl is not None:
        raise CUDAError(f"{name}: {expl}")
    desc_err, desc = runtime.cudaGetErrorString(error)
    if desc_err != _RUNTIME_SUCCESS:
        raise CUDAError(f"{name}")
    desc = desc.decode()
    raise CUDAError(f"{name}: {desc}")


cpdef inline int _check_nvrtc_error(error, handle=None) except?-1:
    if error == _NVRTC_SUCCESS:
        return 0
    err = f"{error}: {nvrtc.nvrtcGetErrorString(error)[1].decode()}"
    if handle is not None:
        _, logsize = nvrtc.nvrtcGetProgramLogSize(handle)
        log = b" " * logsize
        _ = nvrtc.nvrtcGetProgramLog(handle, log)
        err += f", compilation log:\n\n{log.decode('utf-8', errors='backslashreplace')}"
    raise NVRTCError(err)


cdef inline int _check_error(error, handle=None) except?-1:
    if isinstance(error, driver.CUresult):
        return _check_driver_error(error)
    elif isinstance(error, runtime.cudaError_t):
        return _check_runtime_error(error)
    elif isinstance(error, nvrtc.nvrtcResult):
        return _check_nvrtc_error(error, handle=handle)
    else:
        raise RuntimeError(f"Unknown error type: {error}")


def handle_return(result: tuple[Any, ...], handle: object = None) -> Any:
    _check_error(result[0], handle=handle)
    cdef int out_len = len(result)
    if out_len == 1:
        return
    elif out_len == 2:
        return result[1]
    else:
        return result[1:]


cpdef object check_or_create_options(type cls, object options, str options_description="", bint keep_none=False):
    """
    Create the specified options dataclass from a dictionary of options or None.
    """
    if options is None:
        if keep_none:
            return options
        return cls()
    elif isinstance(options, cls):
        return options
    elif isinstance(options, dict):
        return cls(**options)
    else:
        raise TypeError(
            f"The {options_description} must be provided as an object "
            f"of type {cls.__name__} or as a dict with valid {options_description}. "
            f"The provided object is '{options}'."
        )


def _handle_boolean_option(option: bool) -> str:
    """
    Convert a boolean option to a string representation.
    """
    return "true" if bool(option) else "false"


def precondition(checker: Callable[..., None], str what="") -> Callable[..., Any]:
    """
    A decorator that adds checks to ensure any preconditions are met.

    Args:
        checker: The function to call to check whether the preconditions are met. It has
        the same signature as the wrapped function with the addition of the keyword argument `what`.
        what: A string that is passed in to `checker` to provide context information.

    Returns:
        Callable: A decorator that creates the wrapping.
    """

    def outer(wrapped_function: Callable) -> Callable:
        """
        A decorator that actually wraps the function for checking preconditions.
        """

        @functools.wraps(wrapped_function)
        def inner(*args, **kwargs) -> Any:
            """
            Check preconditions and if they are met, call the wrapped function.
            """
            checker(*args, **kwargs, what=what)
            result = wrapped_function(*args, **kwargs)

            return result

        return inner

    return outer


def is_sequence(obj: object) -> bool:
    """
    Check if the given object is a sequence (list or tuple).
    """
    return isinstance(obj, Sequence)


def is_nested_sequence(obj: object) -> bool:
    """
    Check if the given object is a nested sequence (list or tuple with atleast one list or tuple element).
    """
    return is_sequence(obj) and any(is_sequence(elem) for elem in obj)



class Transaction:
    """
    A context manager for transactional operations with undo capability.

    The Transaction class allows you to register undo actions (callbacks) that will be executed
    if the transaction is not committed before exiting the context. This is useful for managing
    resources or operations that need to be rolled back in case of errors or early exits.

    Usage:
        with Transaction() as txn:
            txn.append(some_cleanup_function, arg1, arg2)
            # ... perform operations ...
            txn.commit()  # Disarm undo actions; nothing will be rolled back on exit

    Methods:
        append(fn, *args, **kwargs): Register an undo action to be called on rollback.
        commit(): Disarm all undo actions; nothing will be rolled back on exit.
    """
    def __init__(self) -> None:
        self._stack = ExitStack()
        self._entered = False

    def __enter__(self):
        self._stack.__enter__()
        self._entered = True
        return self

    def __exit__(self, exc_type, exc, tb):
        # If exit callbacks remain, they'll run in LIFO order.
        self._entered = False
        return self._stack.__exit__(exc_type, exc, tb)

    def append(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None:
        """
        Register an undo action (runs if the with-block exits without commit()).
        Values are bound now via partial so late mutations don't bite you.
        """
        if not self._entered:
            raise RuntimeError("Transaction must be entered before append()")
        self._stack.callback(partial(fn, *args, **kwargs))

    def commit(self) -> None:
        """
        Disarm all undo actions. After this, exiting the with-block does nothing.
        """
        # pop_all() empties this stack so no callbacks are triggered on exit.
        self._stack.pop_all()


# Track whether we've already warned about fork method
_fork_warning_checked = False


def reset_fork_warning() -> None:
    """Reset the fork warning check flag for testing purposes.

    This function is intended for use in tests to allow multiple test runs
    to check the warning behavior.
    """
    global _fork_warning_checked
    _fork_warning_checked = False


cdef inline tuple _read_fill_ptr(const char* ptr, Py_ssize_t width):
    """Extract (value, element_size) from a raw pointer of known width."""
    cdef unsigned int val
    if width == 1:
        val = (<uint8_t*>ptr)[0]
    elif width == 2:
        val = (<uint16_t*>ptr)[0]
    elif width == 4:
        val = (<uint32_t*>ptr)[0]
    else:
        raise ValueError(f"value must be 1, 2, or 4 bytes, got {width}")
    return (val, <unsigned int>width)


cpdef tuple _parse_fill_value(value):
    """Parse a fill/memset value into (raw_value, element_size).

    Parameters
    ----------
    value : int or buffer-protocol object
        - int: Must be in range [0, 256). Treated as 1-byte fill.
        - bytes or buffer-protocol: Must be 1, 2, or 4 bytes.

    Returns
    -------
    tuple of (int, int)
        (raw_value, element_size) where element_size is 1, 2, or 4.

    Raises
    ------
    OverflowError
        If int value is outside [0, 256).
    TypeError
        If value is not an int and does not support the buffer protocol.
    ValueError
        If value byte length is not 1, 2, or 4.
    """
    cdef uint8_t byte_val
    cdef Py_buffer buf

    if isinstance(value, int):
        byte_val = value
        return (<unsigned int>byte_val, <unsigned int>1)

    if isinstance(value, bytes):
        return _read_fill_ptr(<const char*><bytes>value, <Py_ssize_t>len(value))

    if PyObject_GetBuffer(value, &buf, PyBUF_SIMPLE) != 0:
        raise TypeError(
            f"value must be an int or support the buffer protocol, "
            f"got {type(value).__name__}"
        )
    try:
        return _read_fill_ptr(<const char*>buf.buf, buf.len)
    finally:
        PyBuffer_Release(&buf)


def check_multiprocessing_start_method() -> None:
    """Check if multiprocessing start method is 'fork' and warn if so."""
    global _fork_warning_checked
    if _fork_warning_checked:
        return
    _fork_warning_checked = True

    # Common warning message parts
    common_message = (
        "CUDA does not support. Forked subprocesses exhibit undefined behavior, "
        "including failure to initialize CUDA contexts and devices. Set the start method "
        "to 'spawn' before creating processes that use CUDA. "
        "Use: multiprocessing.set_start_method('spawn')"
    )

    try:
        start_method = multiprocessing.get_start_method()
        if start_method == "fork":
            message = f"multiprocessing start method is 'fork', which {common_message}"
            warnings.warn(message, UserWarning, stacklevel=3)
    except RuntimeError:
        # get_start_method() can raise RuntimeError if start method hasn't been set
        # In this case, default is 'fork' on Linux, so we should warn
        if platform.system() == "Linux":
            message = (
                f"multiprocessing start method is not set and defaults to 'fork' on Linux, "
                f"which {common_message}"
            )
            warnings.warn(message, UserWarning, stacklevel=3)
