# SPDX-FileCopyrightText: Copyright (c) 2020-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://nvidia.github.io/NVTX/LICENSE.txt for license information.

import contextlib
import os
import warnings

from functools import wraps, lru_cache
from typing import Optional, Union, Tuple

try:
    import numpy as np
except ImportError:
    np = None

from nvtx._lib import (
    Domain,
    DummyDomain,
    dummy_domain,
    mark as libnvtx_mark,
    pop_range as libnvtx_pop_range,
    push_range as libnvtx_push_range,
    start_range as libnvtx_start_range,
    end_range as libnvtx_end_range,
    EntryKind,
    PayloadEntryType,
    NvtxWarning,
)

from nvtx._lib.counters import CounterSemantics
from nvtx._dtype_validation import _validate_entry_storage
from nvtx._metadata import NVTX_DTYPE_METADATA_KEY_NAME, _PayloadMetadata


PayloadTypeAlias = Union[int, float, list, tuple, range, bytes, "np.ndarray"]
_immutable_payload_types = {int, float, tuple, range, bytes}
_ENABLED = not os.getenv("NVTX_DISABLE", False)


def numpy_dtype(
    *args,
    counter_semantics: Optional[CounterSemantics] = None,
    entry_kind: Optional[EntryKind] = None,
    entry_type: Optional[PayloadEntryType] = None,
    **kwargs,
) -> "np.dtype":
    """
    Construct a NumPy dtype, optionally carrying NVTX payload metadata.

    This function accepts the same arguments as :func:`numpy.dtype`. When
    ``counter_semantics`` is provided, the returned dtype carries metadata
    for use as a flat structured counter field. For a top-level scalar counter,
    pass semantics to :meth:`nvtx.Domain.get_counter` instead.

    When ``entry_kind`` is provided, the dtype carries that field role for a
    range/mark field (or a counter-group timestamp).

    Parameters
    ----------
    *args
        Positional arguments passed to :func:`numpy.dtype`.
    counter_semantics : CounterSemantics, optional
        Semantics to attach to the dtype for use as a payload schema entry,
        typically a field in a structured counter dtype. If the NVTX metadata
        key is already present in ``metadata``, the existing value is preserved
        and a warning is emitted.
    entry_kind : EntryKind, optional
        Role of the field in a range/mark or counter-group schema.
        Timestamp roles require 64-bit integer storage.
    entry_type : PayloadEntryType, optional
        Special NVTX interpretation of the field's storage. This distinguishes
        range IDs, category IDs, packed ARGB colors, scope IDs, and
        :attr:`PayloadEntryType.REGISTERED_STRING` handles from ordinary
        integers with the same storage width.
    **kwargs
        Keyword arguments passed to :func:`numpy.dtype`.

    Returns
    -------
    numpy.dtype
        The constructed dtype.

    Raises
    ------
    TypeError
        If an entry role or type is invalid or incompatible with the dtype's
        storage.
    RuntimeError
        If NumPy is not installed.
    """
    if np is None:
        raise RuntimeError("Install numpy to construct NVTX NumPy dtypes.")
    if entry_kind is not None and not isinstance(entry_kind, EntryKind):
        raise TypeError("entry_kind must be an nvtx.EntryKind or None")
    if entry_type is not None and not isinstance(entry_type, PayloadEntryType):
        raise TypeError("entry_type must be an nvtx.PayloadEntryType or None")

    base_dtype = np.dtype(*args, **kwargs)
    _validate_entry_storage(base_dtype, entry_kind)
    _validate_entry_storage(base_dtype, entry_type)
    if counter_semantics is None and entry_kind is None and entry_type is None:
        return base_dtype

    metadata = dict(base_dtype.metadata or {})
    if NVTX_DTYPE_METADATA_KEY_NAME in metadata:
        warnings.warn(
            "NVTX dtype metadata key already exists; leaving it unchanged.",
            NvtxWarning,
            stacklevel=2,
        )
        return base_dtype

    metadata[NVTX_DTYPE_METADATA_KEY_NAME] = _PayloadMetadata(
        counter_semantics=counter_semantics,
        entry_kind=entry_kind,
        entry_type=entry_type,
    )
    return np.dtype(base_dtype, metadata=metadata)


@lru_cache(maxsize=None)
def _get_domain_cached(name):
    return Domain(name)


def get_domain(name: Optional[str] = None) -> Union[Domain, DummyDomain]:
    """
    Get or create a :class:`Domain` object for a domain name.
    """
    return _get_domain_cached(name)


class annotate:
    """
    Annotate code ranges using a context manager or a decorator.

    Parameters
    ----------
    message
        A message associated with the annotated code range.
        When used as a decorator, defaults to the decorated function name.
        When used as a context manager, defaults to the empty string.
        Messages are cached and are registered as Registered Strings
        in NVTX.
        Caching a very large number of messages may lead to increased
        memory usage.
    color
        A color associated with the annotated code range.
        Supports `matplotlib` colors if it is available.
    domain
        A string specifying the domain under which the code range is
        scoped. The default domain is named "NVTX".
    category
        A string or an integer specifying the category within the domain
        under which the code range is scoped. If unspecified, the code
        range is not associated with a category.
    payload
        A value associated with this event.
        Using payloads provides a separation between the message and the data of the event,
        which is often useful for analysis.

        .. note:: payloads of type other than ``int`` or ``float`` requires
                  NumPy to be installed (not installed with ``nvtx`` package).

    Examples
    --------
    >>> import nvtx
    >>> import time

    Using a decorator:

    >>> @nvtx.annotate(color="red", domain="my_domain"):    # `message` defaults to "func"
    ... def func():
    ...     time.sleep(0.1)

    Using a context manager:

    >>> with nvtx.annotate("my_code_range", color="blue"):
    ...    time.sleep(10)
    """

    def __init__(
        self,
        message: Optional[str] = None,
        color: Optional[Union[str, int]] = None,
        domain: Optional[str] = None,
        category: Optional[Union[str, int]] = None,
        payload: Optional[PayloadTypeAlias] = None,
    ):

        self.init_args = message, color, domain, category, payload
        self.domain = get_domain(domain)
        if self.domain is not dummy_domain:
            if payload is not None and type(payload) not in _immutable_payload_types:
                self.mutable_payload = payload
                self.attributes = self.domain.get_event_attributes(
                    message, color, category
                )
            else:
                self.mutable_payload = None
                self.attributes = self.domain.get_event_attributes(
                    message, color, category, payload
                )

    def __reduce__(self):
        return self.__class__, self.init_args

    def __enter__(self):
        if self.domain is not dummy_domain:
            if self.mutable_payload is not None:
                self.attributes.payload = self.mutable_payload
            libnvtx_push_range(self.attributes, self.domain.handle)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self.domain is not dummy_domain:
            libnvtx_pop_range(self.domain.handle)
        return False

    def __call__(self, func):
        if self.domain is dummy_domain:
            return func
        if self.attributes.message is None:
            self.attributes.message = self.domain.get_registered_string(func.__name__)

        @wraps(func)
        def inner(*args, **kwargs):
            if self.mutable_payload is not None:
                self.attributes.payload = self.mutable_payload
            libnvtx_push_range(self.attributes, self.domain.handle)
            try:
                result = func(*args, **kwargs)
            finally:
                # Always pop the range, even if an exception is raised
                libnvtx_pop_range(self.domain.handle)
            return result

        return inner


def mark(
    message: Optional[str] = None,
    color: Optional[Union[str, int]] = "blue",
    domain: Optional[str] = None,
    category: Optional[Union[str, int]] = None,
    payload: Optional[PayloadTypeAlias] = None,
):
    """
    Mark an instantaneous event.

    Parameters
    ----------
    message
        A message associated with the event.
        Messages are cached and are registered as Registered Strings
        in NVTX.
        Caching a very large number of messages may lead to increased
        memory usage.
    color
        Color associated with the event.
        Supports `matplotlib` colors if it is available.
    domain
        A string specifying the domain under which the event is scoped.
        The default domain is named "NVTX".
    category
        A string or an integer specifying the category within the domain
        under which the event is scoped. If unspecified, the event is
        not associated with a category.
    payload
        A value associated with this event.
        Using payloads provides a separation between the message and the data of the event,
        which is often useful for analysis.

        .. note:: payloads of type other than ``int`` or ``float`` requires
                  NumPy to be installed (not installed with ``nvtx`` package).

    Notes
    -----
    For best performance, use :func:`Domain.mark`.

    Examples
    --------
    >>> import nvtx
    >>> nvtx.mark("my_marker", domain="my_domain")
    """
    domain = get_domain(domain)
    if domain is not dummy_domain:
        libnvtx_mark(
            domain.get_event_attributes(message, color, category, payload),
            domain.handle,
        )


def push_range(
    message: Optional[str] = None,
    color: Optional[Union[str, int]] = "blue",
    domain: Optional[str] = None,
    category: Optional[Union[str, int]] = None,
    payload: Optional[PayloadTypeAlias] = None,
):
    """
    Mark the beginning of a code range.

    Parameters
    ----------
    message
        A message associated with the annotated code range.
        Messages are cached and are registered as Registered Strings
        in NVTX.
        Caching a very large number of messages may lead to increased
        memory usage.
    color
        A color associated with the annotated code range.
        Supports `matplotlib` colors if it is available.
    domain
        Name of a domain under which the code range is scoped.
        The default domain name is "NVTX".
    category
        A string or an integer specifying the category within the domain
        under which the code range is scoped. If unspecified, the code range
        is not associated with a category.
    payload
        A value associated with this event.
        Using payloads provides a separation between the message and the data of the event,
        which is often useful for analysis.

        .. note:: payloads of type other than ``int`` or ``float`` requires
                  NumPy to be installed (not installed with ``nvtx`` package).

    Notes
    -----
    When applicable, prefer to use :class:`annotate`.
    Otherwise, for best performance, use :func:`Domain.push_range` and :func:`Domain.pop_range`.

    Examples
    --------
    >>> import time
    >>> import nvtx
    >>> nvtx.push_range("my_code_range", domain="my_domain")
    >>> time.sleep(1)
    >>> nvtx.pop_range(domain="my_domain")
    """
    domain = get_domain(domain)
    if domain is not dummy_domain:
        libnvtx_push_range(
            domain.get_event_attributes(message, color, category, payload),
            domain.handle,
        )


def pop_range(domain: Optional[str] = None):
    """
    Mark the end of a code range that was started with :func:`push_range`.

    Parameters
    ----------
    domain
        The domain under which the code range is scoped. The default
        domain is "NVTX".
    """
    domain = get_domain(domain)
    if domain is not dummy_domain:
        libnvtx_pop_range(domain.handle)


def start_range(
    message: Optional[str] = None,
    color: Optional[Union[str, int]] = None,
    domain: Optional[str] = None,
    category: Optional[Union[str, int]] = None,
    payload: Optional[PayloadTypeAlias] = None,
) -> Tuple[int, int]:
    """
    Mark the beginning of a process range.

    Parameters
    ----------
    message
        A message associated with the range.
        Messages are cached and are registered as Registered Strings
        in NVTX.
        Caching a very large number of messages may lead to increased
        memory usage.
    color
        A color associated with the range.
        Supports `matplotlib` colors if it is available.
    domain
        Name of a domain under which the range is scoped.
        The default domain name is "NVTX".
    category
        A string or an integer specifying the category within the domain
        under which the range is scoped. If unspecified, the range
        is not associated with a category.
    payload
        A value associated with this event.
        Using payloads provides a separation between the message and the data of the event,
        which is often useful for analysis.

        .. note:: payloads of type other than ``int`` or ``float`` requires
                  NumPy to be installed (not installed with ``nvtx`` package).

    Returns
    -------
    A tuple of the form ``(range_id, domain_handle)`` that must be passed to :func:`end_range`.

    Notes
    -----
    For best performance, use :func:`Domain.start_range` and :func:`Domain.end_range`.

    Examples
    --------
    >>> import time
    >>> import nvtx
    >>> range_id = nvtx.start_range("my_code_range", domain="my_domain")
    >>> time.sleep(1)
    >>> nvtx.end_range(range_id)
    """
    domain = get_domain(domain)
    if domain is not dummy_domain:
        return libnvtx_start_range(
            domain.get_event_attributes(message, color, category, payload),
            domain.handle,
        )


def end_range(range_id: Tuple[int, int]):
    """
    Mark the end of a code range that was started with :func:`start_range`.

    Parameters
    ----------
    range_id
        The tuple object returned by :func:`start_range`.
    """
    if range_id is not None:
        libnvtx_end_range(*range_id)


def enabled():
    """
    Returns True if nvtx is enabled.
    """
    return _ENABLED


if not enabled():

    class annotate(contextlib.nullcontext):
        def __init__(self, *args, **kwargs):
            super().__init__()

        def __call__(self, func):
            return func

    # Could use a decorator here but overheads are significant enough
    # not to. See https://github.com/NVIDIA/NVTX/pull/24 for discussion.
    def mark(message=None, color=None, domain=None, category=None, payload=None):
        pass

    def push_range(message=None, color=None, domain=None, category=None, payload=None):
        pass

    def pop_range(domain=None):
        pass

    def start_range(message=None, color=None, domain=None, category=None, payload=None):
        pass

    def end_range(range_id):
        pass
