# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

"""General helpers for this extension."""

from enum import Enum
from typing import (
    Any,
    Optional,
    Union,
)

import omni.graph.core as og
import warp as wp


#   General
# ------------------------------------------------------------------------------


class IntEnum(int, Enum):
    """Base class for integer enumerators with labels."""

    def __new__(cls, value, label):
        obj = int.__new__(cls, value)
        obj._value_ = value
        obj.label = label
        return obj


#   Timer
# ------------------------------------------------------------------------------


class NodeTimer(object):
    """Context wrapping Warp's scoped timer for use with nodes."""

    def __init__(self, name: str, db: Any, active: bool = False) -> None:
        name = "{}:{}".format(db.node.get_prim_path(), name)
        self.timer = wp.ScopedTimer(name, active=active, synchronize=True)

    def __enter__(self) -> None:
        self.timer.__enter__()
        return self

    def __exit__(self, type: Any, value: Any, traceback: Any) -> None:
        self.timer.__exit__(type, value, traceback)


#   Types
# ------------------------------------------------------------------------------

_BaseDType = og.BaseDataType
_AttrRole = og.AttributeRole

# fmt: off

_DATA_TYPES_MAPPING = (
    ("bool"      , (_BaseDType.BOOL  ,  1, _AttrRole.NONE      ), "int8"   ),
    ("color3f"   , (_BaseDType.FLOAT ,  3, _AttrRole.COLOR     ), "vec3"   ),
    ("color4f"   , (_BaseDType.FLOAT ,  4, _AttrRole.COLOR     ), "vec4"   ),
    ("double"    , (_BaseDType.DOUBLE,  1, _AttrRole.NONE      ), "float64"),
    ("float"     , (_BaseDType.FLOAT ,  1, _AttrRole.NONE      ), "float32"),
    ("float2"    , (_BaseDType.FLOAT ,  2, _AttrRole.NONE      ), "vec2"   ),
    ("float3"    , (_BaseDType.FLOAT ,  3, _AttrRole.NONE      ), "vec3"   ),
    ("float4"    , (_BaseDType.FLOAT ,  4, _AttrRole.NONE      ), "vec4"   ),
    ("int"       , (_BaseDType.INT   ,  1, _AttrRole.NONE      ), "int32"  ),
    ("int64"     , (_BaseDType.INT64 ,  1, _AttrRole.NONE      ), "int64"  ),
    ("matrix2d"  , (_BaseDType.DOUBLE,  4, _AttrRole.MATRIX    ), "mat22d" ),
    ("matrix3d"  , (_BaseDType.DOUBLE,  9, _AttrRole.MATRIX    ), "mat33d" ),
    ("matrix4d"  , (_BaseDType.DOUBLE, 16, _AttrRole.MATRIX    ), "mat44d" ),
    ("normal3f"  , (_BaseDType.FLOAT ,  3, _AttrRole.NORMAL    ), "vec3"   ),
    ("point3f"   , (_BaseDType.FLOAT ,  3, _AttrRole.POSITION  ), "vec3"   ),
    ("quatf"     , (_BaseDType.FLOAT ,  4, _AttrRole.QUATERNION), "quat"   ),
    ("texCoord2f", (_BaseDType.FLOAT ,  2, _AttrRole.TEXCOORD  ), "vec2"   ),
    ("texCoord3f", (_BaseDType.FLOAT ,  3, _AttrRole.TEXCOORD  ), "vec3"   ),
    ("timecode"  , (_BaseDType.DOUBLE,  1, _AttrRole.TIMECODE  ), "float64"),
    ("token"     , (_BaseDType.TOKEN ,  1, _AttrRole.NONE      ), "uint64" ),
    ("uchar"     , (_BaseDType.UCHAR ,  1, _AttrRole.NONE      ), "uint8"  ),
    ("uint"      , (_BaseDType.UINT  ,  1, _AttrRole.NONE      ), "uint32" ),
    ("uint64"    , (_BaseDType.UINT64,  1, _AttrRole.NONE      ), "uint64" ),
    ("vector3f"  , (_BaseDType.FLOAT ,  3, _AttrRole.VECTOR    ), "vec3"   ),
)

_SDF_DATA_TYPE_TO_OG        = {k: v for (k, v, _) in _DATA_TYPES_MAPPING}
_SDF_DATA_TYPE_NAME_TO_WARP = {k: v for (k, _, v) in _DATA_TYPES_MAPPING}
_OG_DATA_TYPE_TO_WARP       = {k: v for (_, k, v) in _DATA_TYPES_MAPPING}

# fmt: on

SUPPORTED_OG_DATA_TYPES = tuple(
    og.Type(base_data_type, tuple_count=tuple_count, array_depth=0, role=role)
    for base_data_type, tuple_count, role in _OG_DATA_TYPE_TO_WARP.keys()
)
SUPPORTED_OG_ARRAY_TYPES = tuple(
    og.Type(base_data_type, tuple_count=tuple_count, array_depth=1, role=role)
    for base_data_type, tuple_count, role in _OG_DATA_TYPE_TO_WARP.keys()
)
SUPPORTED_OG_TYPES = SUPPORTED_OG_DATA_TYPES + SUPPORTED_OG_ARRAY_TYPES

SUPPORTED_SDF_DATA_TYPE_NAMES = tuple(_SDF_DATA_TYPE_NAME_TO_WARP.keys())


def get_warp_type_from_data_type_name(
    data_type_name: str,
    dim_count: int = 0,
    as_str: bool = False,
    str_namespace: Optional[str] = "wp",
):
    if as_str:
        prefix = "" if str_namespace is None else "{}.".format(str_namespace)

        if dim_count == 0:
            return "{prefix}{dtype}".format(prefix=prefix, dtype=data_type_name)

        if dim_count == 1:
            return "{prefix}array(dtype={prefix}{dtype})".format(
                prefix=prefix,
                dtype=data_type_name,
            )

        return "{prefix}array(dtype={prefix}{dtype}, ndim={ndim})".format(
            prefix=prefix,
            dtype=data_type_name,
            ndim=dim_count,
        )

    dtype = getattr(wp.types, data_type_name)
    if dim_count == 0:
        return dtype

    if dim_count == 1:
        return wp.array(dtype=dtype)

    return wp.array(dtype=dtype, ndim=dim_count)


def type_convert_og_to_warp(
    og_type: og.Type,
    dim_count: Optional[int] = None,
    as_str: bool = False,
    str_namespace: Optional[str] = "wp",
) -> Union[Any, str]:
    """Converts an OmniGraph type into a compatible Warp type."""
    data_type_name = _OG_DATA_TYPE_TO_WARP.get(
        (og_type.base_type, og_type.tuple_count, og_type.role),
    )
    if data_type_name is None:
        raise RuntimeError("Unsupported attribute type '{}'.".format(og_type))

    if dim_count is None:
        dim_count = og_type.array_depth

    return get_warp_type_from_data_type_name(
        data_type_name,
        dim_count=dim_count,
        as_str=as_str,
        str_namespace=str_namespace,
    )


def type_convert_sdf_name_to_warp(
    sdf_type_name: str,
    dim_count: Optional[int] = None,
    as_str: bool = False,
    str_namespace: Optional[str] = "wp",
) -> Union[Any, str]:
    """Converts a Sdf type name into a compatible Warp type."""
    if sdf_type_name.endswith("[]"):
        sdf_type_name = sdf_type_name[:-2]
        if dim_count is None:
            dim_count = 1
    elif dim_count is None:
        dim_count = 0

    data_type_name = _SDF_DATA_TYPE_NAME_TO_WARP.get(sdf_type_name)
    if data_type_name is None:
        raise RuntimeError("Unsupported attribute type '{}'.".format(sdf_type_name))

    return get_warp_type_from_data_type_name(
        data_type_name,
        dim_count=dim_count,
        as_str=as_str,
        str_namespace=str_namespace,
    )


def type_convert_sdf_name_to_og(
    sdf_type_name: str,
    is_array: Optional[bool] = None,
) -> og.Type:
    """Converts a Sdf type name into its corresponding OmniGraph type."""
    if sdf_type_name.endswith("[]"):
        sdf_type_name = sdf_type_name[:-2]
        if is_array is None:
            is_array = True
    elif is_array is None:
        is_array = False

    data_type = _SDF_DATA_TYPE_TO_OG.get(sdf_type_name)
    if data_type is None:
        raise RuntimeError("Unsupported attribute type '{}'.".format(sdf_type_name))

    base_data_type, tuple_count, role = data_type
    return og.Type(
        base_data_type,
        tuple_count=tuple_count,
        array_depth=int(is_array),
        role=role,
    )
