# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0

"""
A collection of (internal use) helper functions.
"""

import contextlib
import functools
import time
import typing
from collections.abc import Callable, Sequence

import cuda.bindings.runtime as cbr

try:
    from cuda.core import Device, Event, EventOptions
except ImportError:
    from cuda.core.experimental import Device, Event, EventOptions

from nvmath._internal.layout import is_contiguous_and_dense

from . import formatters, mem_limit, package_wrapper
from .package_ifc import AnyStream, Stream, StreamHolder
from .tensor_ifc import AnyTensor, Tensor, TensorHolder


def infer_object_package(obj: object) -> str:
    """
    Infer the package that defines this object.
    """
    module = obj.__class__.__module__
    return module.split(".")[0]


"""A generic type for Options classes"""
Options = typing.TypeVar("Options")


def check_or_create_options(
    cls: type[Options], options: None | dict[str, typing.Any] | Options, options_description: str, *, keep_none=False
) -> Options | None:
    """
    Create the specified options dataclass from a dictionary of options or None.
    """

    if options is None:
        if keep_none:
            return options
        options = cls()
    elif isinstance(options, dict):
        options = cls(**options)

    if not isinstance(options, cls):
        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}'."
        )

    return options


def check_or_create_one_of_options(clss, options, options_description, *, cls_key="name", default_name=None):
    """
    Create one of the specified options dataclasses by name or from a dictionary of options.
    """

    assert isinstance(clss, tuple)
    assert all(hasattr(cls, "name") and isinstance(cls.name, str) for cls in clss)

    if options is None:
        assert default_name is not None, "Internal error in execution options"
        options = _create_one_of_options_from_name(clss, default_name, options_description, cls_key=cls_key)
    elif isinstance(options, str):
        options = _create_one_of_options_from_name(clss, options, options_description, cls_key=cls_key)
    elif isinstance(options, dict):
        options = _create_one_of_options_from_dict(
            clss,
            options,
            options_description,
            cls_key=cls_key,
            default_name=default_name,
        )
    else:
        if not isinstance(options, clss):
            _raise_invalid_one_of_options(clss, options, options_description, cls_key=cls_key)
    return options


def _create_one_of_options_from_name(clss, cls_name, options_description, *, cls_key="name"):
    try:
        _cls_name = cls_name.lower()
        return next(cls for cls in clss if cls.name == _cls_name)()
    except StopIteration:
        raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None


def _create_one_of_options_from_dict(
    clss,
    options,
    options_description,
    *,
    cls_key="name",
    default_name=None,
):
    cls_name = options.get(cls_key)
    if cls_name is None:
        cls_name = default_name
    if not isinstance(cls_name, str):
        raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key)
    try:
        _cls_name = cls_name.lower()
        cls = next(cls for cls in clss if cls.name == _cls_name)
    except StopIteration:
        raise _raise_invalid_one_of_options(clss, cls_name, options_description, cls_key=cls_key) from None
    return cls(**{key: name for key, name in options.items() if key != cls_key})


def _raise_invalid_one_of_options(clss, options, options_description, *, cls_key="name"):
    accepted_names = ", ".join(f"'{cls.name}'" for cls in clss)
    accepted_types = ", ".join(str(cls) for cls in clss)
    raise ValueError(
        f"The {options_description} must be: \n"
        f"1. an object of one of the following types {accepted_types} or\n"
        f"2. a string with the type name ({accepted_names}) or\n"
        f"3. a dict with a '{cls_key}' key and one of the type names {accepted_names}' as a value, "
        f"and optional options valid for that type. \n"
        f"The provided object is '{options}'."
    )


class device_ctx:
    """
    NOTE, using classic class-based context manager here as it has lower overhead
    than the generator-based context manager.

    Semantics:

    1. The device context manager makes the specified device current from the point of entry
       until the point of exit.

    2. When the context manager exits, the current device is reset to what it was when the
       context manager was entered.

    3. Any explicit setting of the device within the context manager (using
       cupy.cuda.Device().use(), torch.cuda.set_device(), etc) will overrule the device set
       by the context manager from that point onwards till the context manager exits. In
       other words, the context manager provides a local device scope and the current device
       can be explicitly reset for the remainder of that scope.

    4. The context manager is single-use.

    Corollary: if any library function resets the device globally and this is an undesired
        side-effect, such functions must be called from within the device context manager.

    Device context managers can be arbitrarily nested.
    """

    __slots__ = ("new_device_id", "_old_device")

    def __init__(self, new_device_id):
        self.new_device_id = new_device_id
        self._old_device = None

    def __enter__(self):
        if self._old_device is not None:
            raise RuntimeError("Reusing a device_ctx instance is not allowed.")
        self._old_device = old_device = Device()
        new_device_id = self.new_device_id
        if old_device.device_id != new_device_id:
            device = Device(new_device_id)
            device.set_current()
            return device
        else:
            return old_device

    def __exit__(self, type, value, traceback):
        # We should always restore the old device at exit.
        self._old_device.set_current()


def is_hashable(obj: object) -> bool:
    try:
        hash(obj)
    except TypeError:
        return False
    return True


@functools.lru_cache(maxsize=128)
def cached_get_or_create_stream(
    device_id: int,
    stream: Stream | int,
    op_package: str,
) -> StreamHolder[Stream]:
    op_package_ifc = package_wrapper.PACKAGE[op_package]
    if isinstance(stream, int):
        ptr = stream
        external = op_package_ifc.create_external_stream(device_id, ptr)
        ctx = op_package_ifc.to_stream_context(external)
        obj = op_package_ifc.create_stream(external, device_id)
        return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)

    stream_package = infer_object_package(stream)
    if stream_package != op_package:
        message = (
            f"The stream object must belong to the same package as the tensor network operands. "
            f"Stream package: {stream_package}, Tensor package: {op_package}"
        )
        raise TypeError(message)

    ctx = op_package_ifc.to_stream_context(stream)
    ptr = op_package_ifc.to_stream_pointer(stream)
    external = stream
    obj = op_package_ifc.create_stream(stream, device_id)
    return StreamHolder(ctx=ctx, device_id=device_id, external=external, obj=obj, ptr=ptr, package=op_package)


def get_or_create_stream(
    device_id: int,
    stream: Stream | int | None,
    op_package: str,
) -> StreamHolder[Stream]:
    """
    Create a stream object from a stream pointer or extract the stream pointer from a stream
    object, or use the current stream.

    Args:
        device_id: The device ID.

        stream: A stream object, stream pointer, or None.

        op_package: The package the tensor network operands belong to.

    Returns:
        StreamHolder: Hold a CuPy stream object, package stream context, stream pointer, ...
    """
    op_package_ifc = package_wrapper.PACKAGE[op_package]
    if stream is None:
        stream = op_package_ifc.get_current_stream(device_id)
    return cached_get_or_create_stream(device_id, stream, op_package)


@functools.lru_cache(maxsize=128)
def get_memory_limit_from_device_id(memory_limit: int | float | str, device_id: int) -> int:
    with device_ctx(device_id):
        status, _, total_memory = cbr.cudaMemGetInfo()
        if status != 0 or total_memory is None:
            raise RuntimeError(
                f"cudaMemGetInfo failed with status {status}, total_memory={total_memory}. A possible cause is \
an inconsistent version of cuda-bindings for the CTK version used with nvmath-python."
            )
        return _get_memory_limit(memory_limit, total_memory)


def _get_memory_limit(memory_limit: int | float | str, total_memory: int) -> int:
    """
    Parse user provided memory limit and return the memory limit in bytes.
    """
    if isinstance(memory_limit, int):
        if memory_limit < 0:
            raise ValueError("The specified memory limit must be greater than or equal to 0.")
        return memory_limit

    if isinstance(memory_limit, float):
        if memory_limit < 0:
            raise ValueError("The specified memory limit must be greater than or equal to 0.")
        if memory_limit <= 1.0:
            memory_limit *= total_memory
        return int(memory_limit)

    m = mem_limit.MEM_LIMIT_RE_PCT.match(memory_limit)
    if m:
        factor = float(m.group(1))
        if factor <= 0 or factor > 100:
            raise ValueError("The memory limit percentage must be in the range [0, 100].")
        return int(factor * total_memory / 100.0)

    m = mem_limit.MEM_LIMIT_RE_VAL.match(memory_limit)
    if not m:
        raise ValueError(mem_limit.MEM_LIMIT_DOC.format(kind="memory limit", value=memory_limit))

    base = 1000
    if m.group("binary"):
        base = 1024

    powers = {"": 0, "k": 1, "m": 2, "g": 3}
    unit = m.group("units").lower() if m.group("units") else ""
    multiplier = base ** powers[unit]

    value = float(m.group("value"))
    memory_limit = int(value * multiplier)

    return memory_limit


def create_empty_tensor(
    cls: type[TensorHolder[Tensor]],
    extents: Sequence[int],
    dtype: str,
    device_id: int | typing.Literal["cpu"],
    stream_holder: StreamHolder | None,
    verify_strides: bool,
    strides: Sequence[int] | None = None,
    **context,
) -> TensorHolder[Tensor]:
    """
    Create a wrapped tensor of the same type as (the wrapped) cls on the specified device
    having the specified extents and dtype.

    The tensor is created within a stream context to allow for asynchronous memory
    allocators like CuPy's MemoryAsyncPool.

    Note, the function assumes the `strides` are dense (possibly permuted).
    Otherwise, the behaviour is framework specific and tensor creation may fail
    or created tensor may be corrupted. Set `verify_strides` to True to check
    the layout and drop the strides if the layout is not dense.
    """
    assert isinstance(device_id, int) or device_id == "cpu", f"Internal Error: device_id must be int or 'cpu'; not {device_id}"
    # if device id is none the stream holder must be too
    assert device_id != "cpu" or stream_holder is None
    if strides is not None and verify_strides and not is_contiguous_and_dense(extents, strides):
        strides = None
    tensor = cls.empty(extents, dtype=dtype, device_id=device_id, strides=strides, stream_holder=stream_holder, **context)
    return tensor


def get_operands_device_id(operands: Sequence[TensorHolder[AnyTensor]]):
    """
    Return the id (ordinal) of the device the operands are on, or None if it is on the CPU.
    """
    device_id = operands[0].device_id
    if not all(operand.device_id == device_id for operand in operands):
        devices = {operand.device_id for operand in operands}
        raise ValueError(f"All operands are not on the same device. Devices = {devices}.")

    return device_id


def get_operands_dtype(operands: Sequence[TensorHolder[AnyTensor]]):
    """
    Return the data type name of the tensors.
    """
    dtype = operands[0].dtype
    if not all(operand.dtype == dtype for operand in operands[1:]):
        dtypes = {operand.dtype for operand in operands}
        raise ValueError(f"All tensors must have the same data type. Data types found = {dtypes}.")
    return dtype


def get_operands_package(operands: Sequence[TensorHolder[AnyTensor]]) -> str:
    """
    Return the package name of the tensors.
    """
    package = infer_object_package(operands[0].tensor)
    if not all(infer_object_package(operand.tensor) == package for operand in operands[1:]):
        packages = {infer_object_package(operand.tensor) for operand in operands}
        raise TypeError(f"All tensors must be from the same library package. Packages found = {packages}.")
    return package


def check_operands_match(
    orig_operands: Sequence[object] | object, new_operands: Sequence[object] | object, attribute: str, description: str
):
    """
    Check if the specified attribute matches between the corresponding new and old operands,
    and raise an exception if it doesn't.
    """
    if isinstance(orig_operands, Sequence):
        assert isinstance(new_operands, Sequence), "Internal Error. Cannot compare operand sequence with non-sequence."
        checks: list[bool] = [
            getattr(o, attribute) == getattr(n, attribute) for o, n in zip(orig_operands, new_operands, strict=True)
        ]

        if not all(checks):
            mismatch = [
                f"{location}: {getattr(orig_operands[location], attribute)} => {getattr(new_operands[location], attribute)}"
                for location, predicate in enumerate(checks)
                if predicate is False
            ]
            mismatch = formatters.array2string(mismatch)
            message = f"""\
The {description} of each new operand must match the {description} of the corresponding original operand.

The mismatch in {description} as a sequence of "position: original {description} => new {description}" is: \n{mismatch}"""
            raise ValueError(message)
    else:
        assert isinstance(new_operands, object)
        check: bool = getattr(orig_operands, attribute) == getattr(new_operands, attribute)
        if not check:
            message = f"""The {description} of the new operand must match the {description} of the original operand."""
            raise ValueError(message)


def check_attribute_match(orig_attribute: object, new_attribute: object, description: str):
    """
    Check if the specified attribute matches between the corresponding new and old operands,
    and raise an exception if it doesn't.
    """
    check = orig_attribute == new_attribute
    if not check:
        message = f"""The {description} of the new operand must match the {description} of the original operand."""
        raise ValueError(message)


def get_ptr_from_memory_pointer(mem_ptr):
    """
    Access the value associated with one of the attributes 'device_ptr', 'device_pointer',
    'ptr'.

    This is a compatibility function so that external memory pointer implementations from
    numba and cupy can be used instead of MemoryPointer from nvmath.memory.
    """
    attributes = (
        "device_ptr",  # nvmath.memory
        "device_pointer",
        "ptr",  # cupy.cuda.MemoryPointer
        "pointer",  # numba.cuda.MemoryPointer
    )
    for attr in attributes:
        if hasattr(mem_ptr, attr):
            return getattr(mem_ptr, attr)

    message = f"Memory pointer objects should have one of the following attributes specifying the device pointer: {attributes}"
    raise AttributeError(message)


class Value:
    """
    A simple value wrapper holding a default value.
    """

    def __init__(self, default, *, validator: Callable[[object], bool]):
        """
        Args:
            default: The default value to use.
            validator: A callable that validates the provided value.
        """
        self.validator = validator
        self._data = default

    @property
    def data(self):
        return self._data

    @data.setter
    def data(self, value):
        self._data = self._validate(value)

    def _validate(self, value):
        if self.validator(value):
            return value
        raise ValueError(f"Internal Error: value '{value}' is not valid.")


@contextlib.contextmanager
def cuda_call_ctx(stream_holder: StreamHolder[AnyStream], blocking=True, timing=True) -> typing.Iterator[tuple[Event, Value]]:
    """
    A simple context manager that provides (non-)blocking behavior depending on the
    `blocking` parameter for CUDA calls. The call is timed only for blocking behavior when
    timing is requested.

    An `end` event is recorded after the CUDA call for use in establishing stream ordering
    for non-blocking calls. This event is returned together with a `Value` object that
    stores the elapsed time if the call is blocking and timing is requested, or None
    otherwise.

    Inside the context, the current device is the device which owns the provided stream.
    """
    device_id = stream_holder.device_id
    stream = stream_holder.obj

    with device_ctx(device_id):
        if timing:
            start = stream.record(options=EventOptions(enable_timing=blocking))

        end = stream.device.create_event(options=EventOptions(enable_timing=(timing and blocking)))

        time = Value(None, validator=lambda v: True)
        yield end, time

        stream.record(end)

        if not blocking:
            return

        end.sync()

        if timing:
            time.data = end - start


@contextlib.contextmanager
def host_call_ctx(timing=False):
    elapsed = Value(None, validator=lambda v: True)

    if timing:
        start_time = time.perf_counter_ns()

    yield elapsed

    if timing:
        elapsed.data = (time.perf_counter_ns() - start_time) * 1e-6


# Decorator definitions


def atomic(
    handler: Callable[[typing.Any, Exception | None], bool] | Callable[[Exception | None], bool], method: bool = False
) -> Callable:
    """
    A decorator that provides "succeed or roll-back" semantics. A typical use for this is to
    release partial resources if an exception occurs.

    Args:
        handler: A function to call when an exception occurs. The handler takes a single
            argument, which is the exception object, and returns a boolean stating whether
            the same exception should be reraised. We assume that this function does not
            raise an exception.

        method: Specify if the wrapped function as well as the exception handler are methods
            bound to the same object (method = True) or they are free functions (method =
            False).

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

    def outer(wrapped_function):
        """
        A decorator that actually wraps the function for exception handling.
        """

        @functools.wraps(wrapped_function)
        def inner(*args, **kwargs):
            """
            Call the wrapped function and return the result. If an exception occurs, then
            call the exception handler and reraise the exception.
            """
            try:
                result = wrapped_function(*args, **kwargs)
            except BaseException as e:
                if method:
                    flag = handler(args[0], e)
                else:
                    flag = handler(e)

                if flag:
                    raise e

            return result

        return inner

    return outer


def precondition(checker: Callable[..., None], what: str = "") -> Callable:
    """
    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):
        """
        A decorator that actually wraps the function for checking preconditions.
        """

        @functools.wraps(wrapped_function)
        def inner(*args, **kwargs):
            """
            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


COMMON_SHARED_DOC_MAP = {
    "operand": """\
A tensor (ndarray-like object). The currently supported types are :class:`numpy.ndarray`,
:class:`cupy.ndarray`, and :class:`torch.Tensor`.""".replace("\n", " "),
    #
    "stream": """\
Provide the CUDA stream to use for executing the operation. Acceptable inputs include
``cudaStream_t`` (as Python :class:`int`), :class:`cupy.cuda.Stream`, and
:class:`torch.cuda.Stream`. If a stream is not provided, the current stream from the operand
package will be used. See :ref:`stream-semantics-guide` for more details on stream
handling.""".replace("\n", " "),
    #
    "release_workspace": """\
A value of `True` specifies that the stateful object should release workspace memory back to
the package memory pool on function return, while a value of `False` specifies that the
object should retain the memory. This option may be set to `True` if the application
performs other operations that consume a lot of memory between successive calls to the (same
or different) :meth:`execute` API, but incurs a small overhead due to obtaining and
releasing workspace memory from and to the package memory pool on every call. The default is
`False`.""".replace("\n", " "),
}


def _release_operand_docstring(plural, version_added=None):
    operands = "operands" if plural else "operand"
    their = "their" if plural else "its"
    counts = "counts" if plural else "count"
    reside = "reside" if plural else "resides"
    reset_method = "reset_operands" if plural else "reset_operand"
    reset_method_unchecked = "reset_operands_unchecked" if plural else "reset_operand_unchecked"
    new_operands = "new operands" if plural else "a new operand"
    are = "are" if plural else "is"
    version_line = f"\n\n        .. versionadded:: {version_added}" if version_added else ""
    return f"""\
.. experimental:: method
{version_line}

        This method does two things:

        - Releases internal references to the user-provided {operands},
          so that this instance no longer contributes to {their} reference {counts}.

        - Frees any internal copies (mirrors) that were created when the
          user-provided {operands} {reside} in a different memory space than
          the execution (i.e., copies made during construction or
          :meth:`{reset_method}` / :meth:`{reset_method_unchecked}`
          if present).

        This functionality can be useful in memory-constrained scenarios, e.g. where
        multiple stateful objects need to coexist. Leveraging this functionality,
        the caller can reduce memory usage while retaining the planned state.

        Args:
            None

        Returns:
            None

        Semantics:
            - Preserves the planned state of the stateful object.

            - After calling this method, :meth:`{reset_method}` (or
              :meth:`{reset_method_unchecked}` if present) must be called to
              supply {new_operands} before the next :meth:`execute` call.
              Failure to do so will result in a runtime error.
              Device-side copies will be re-allocated as needed.

            - For cross-space scenarios (e.g. CPU {operands} with GPU execution,
              or GPU {operands} with CPU execution): execution is guaranteed to be
              always blocking, so :meth:`execute` does not return until all
              computation is complete. It is therefore always safe to call this
              method after calling :meth:`execute` without additional synchronization.

            - When the {operands} {are} in the same memory space as the execution
              (e.g. GPU {operands} with GPU execution): in such case, this method
              drops this instance's internal reference to the user-provided {operands}.
              If the reference count of the {operands} reaches zero, {their} memory
              may be freed, so particular attention should be paid. The caller is
              responsible to ensure that if such deallocation happens, it is
              ordered after pending computation (e.g. by retaining a reference until
              the computation is complete, or by synchronizing the stream).
              Failure to do so is analogous to use-after-free.

            See :ref:`nvmath overview`, :ref:`stateful_apis_guide` for operand
            lifecycle and usage patterns, and :ref:`stream-semantics-guide`
            for stream ordering rules.
"""


COMMON_SHARED_DOC_MAP["release_operands"] = _release_operand_docstring(True, version_added="0.9.0")
COMMON_SHARED_DOC_MAP["release_operand"] = _release_operand_docstring(False, version_added="0.9.0")


def _reset_operand_unchecked_docstring(plural, version_added=None, validation_examples="package match, data type match"):
    operands = "operands" if plural else "operand"
    reset_method = "reset_operands" if plural else "reset_operand"
    release_method = "release_operands" if plural else "release_operand"
    version_line = f"\n\n        .. versionadded:: {version_added}" if version_added else ""
    return f"""\
.. experimental:: method
{version_line}

        This method is a performance-optimized alternative to :meth:`{reset_method}`
        that eliminates validation and logging overhead, making it ideal for
        performance-critical loops where {operands}
        compatibility is guaranteed by the caller.

        This method accepts the same parameters as :meth:`{reset_method}`.

        Semantics:
            The semantics are the same as in :meth:`{reset_method}`,
            except that this method does not perform any validation
            (e.g. {validation_examples}, etc.) or logging.

        When to Use:
            - Performance-critical loops with repeated executions on different {operands}

            - After verifying correctness with :meth:`{reset_method}` during development

            - When {operands} compatibility is guaranteed by construction or invariant

        .. seealso::
            :meth:`{reset_method}`, :meth:`{release_method}`
"""


COMMON_SHARED_DOC_MAP["reset_operands_unchecked"] = _reset_operand_unchecked_docstring(True, version_added="0.9.0")
COMMON_SHARED_DOC_MAP["reset_operand_unchecked"] = _reset_operand_unchecked_docstring(False, version_added="0.9.0")


class DefaultDocstring(dict):
    def __missing__(self, key):
        return "{" + key + "}"


def docstring_decorator(doc_map, skip_missing=False):
    assert isinstance(doc_map, dict)

    def _format_doc(doc):
        if doc is not None:
            if skip_missing:
                # Using a default class to handle missing keys in doc_map
                doc = doc.format_map(DefaultDocstring(**doc_map))
            else:
                doc = doc.format(**doc_map)
        return doc

    def decorator(func_or_class):
        if isinstance(func_or_class, type):  # class decorator
            # update the docstring of all public methods with docstrings
            static_methods = []  # staticmethods appear to require special handling
            for name, method in vars(func_or_class).items():
                if isinstance(method, (staticmethod, classmethod)):
                    static_methods.append(name)
                    continue
                if isinstance(method, property):
                    if method.fget and method.fget.__doc__:
                        # Create a new property with updated docstring
                        setattr(
                            func_or_class,
                            name,
                            property(method.fget, method.fset, method.fdel, _format_doc(method.fget.__doc__)),
                        )
                    continue
                if callable(method) and (not name.startswith("_")) and method.__doc__:
                    method.__doc__ = _format_doc(method.__doc__)
            # update the docstring of the constructor
            func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
            for name in static_methods:
                method = func_or_class.__dict__[name]
                method.__func__.__doc__ = _format_doc(method.__func__.__doc__)
            return func_or_class
        else:  # function decorator
            func_or_class.__doc__ = _format_doc(func_or_class.__doc__)
            return func_or_class

    return decorator
