// Copyright (c) 2020, 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.
//
#pragma once

#include <sstream>
#include <string>

namespace carb
{
namespace flatcache
{

// TypeC are integer keys that identify types, like float3, int[] etc.
// There isn't a USD type that can be cast to TypeC,
// please use carb::flatcache::usdTypeToTypeC().
struct TypeC
{
    uint32_t type;
    constexpr bool operator<(const TypeC& other) const
    {
        return type < other.type;
    }
    constexpr bool operator==(const TypeC& other) const
    {
        return type == other.type;
    }
    constexpr bool operator!=(const TypeC& other) const
    {
        return type != other.type;
    }
};

static_assert(std::is_standard_layout<TypeC>::value, "Struct must be standard layout as it is used in C-ABI interfaces");

static constexpr TypeC kUnknownType{ 0 };

enum class BaseDataType : uint8_t
{
    eUnknown = 0,
    eBool,
    eUChar,
    eInt,
    eUInt,
    eInt64,
    eUInt64,

    eHalf,
    eFloat,
    eDouble,

    eToken,

    // RELATIONSHIP is stored as a 64-bit integer internally, but shouldn't be
    // treated as an integer type by nodes.
    eRelationship,

    // For internal use only
    eAsset,
    ePrim,
    eConnection,
    // eTags are attributes that have a name but no type or value
    // They are used for named tags, including USD applied schemas
    eTag
};

inline std::ostream& operator<<(std::ostream& s, const BaseDataType& type)
{
    static const std::string names[] = { "unknown", "bool",   "uchar", "int",   "uint",
                                         "int64",   "uint64", "half",  "float", "double",
                                         "token",   "rel",    "asset", "prim",  "connection",
                                         "tag" };
    if (type <= BaseDataType::eTag)
    {
        return s << names[uint8_t(type)];
    }
    return s;
}

// These correspond with USD attribute "roles", with the exception of eString.
// For example that a vector3f or vector3d (VECTOR) would be transformed
// differently from a point3f or point3d (POSITION).
enum class AttributeRole : uint8_t
{
    eNone = 0,
    eVector,
    eNormal,
    ePosition,
    eColor,
    eTexCoord,
    eQuaternion,
    eTransform,
    eFrame,
    eTimeCode,
    // eText is not a USD role. If a uchar[] attribute has role eText then
    // the corresponding USD attribute will have type "string", and be human
    // readable in USDA. If it doesn't, then it will have type "uchar[]" in USD
    // and appear as an array of numbers in USDA.
    eText,
    // eAppliedSchema is not a USD role, eTags with this role are USD applied schema.
    eAppliedSchema,
    // ePrimTypeName is not a USD role, eTags with this role are USD prim types.
    ePrimTypeName,
    // eExecution is not a USD role, uint attributes with this role are used for control flow in Action Graphs.
    eExecution,
    eMatrix,
    // eObjectId is not a USD role, uint64 attributes with this role are used for Python object identification.
    eObjectId,
    // eBundle is not a USD role, ePrim and eRelationship attributes with this role identify OmniGraph bundles
    eBundle,
    // ePath is not a USD role, it refers to strings that are reinterpreted as SdfPaths. The attribute type must
    // be a uchar[] with a USD type "string".
    ePath,
    // eInstancedAttribute is used as a role on tag types in place of attribute types on instanced prims.
    eInstancedAttribute,
    // eAncestorPrimTypeName is not a USD role, eTags with this role are ancestor types of a USD prim type.
    eAncestorPrimTypeName,
    // Special marker for roles that are not yet determined
    eUnknown,
};

inline std::ostream& operator<<(std::ostream& s, const AttributeRole& type)
{
    static const std::string ognRoleNames[] = { "none",     "vector",        "normal",       "point",              "color",
                                                "texcoord", "quat",          "transform",    "frame",              "timecode",
                                                "text",     "appliedSchema", "primTypeName", "execution",          "matrix",
                                                "objectId", "bundle",        "path",         "instancedAttribute", "ancestorPrimTypeName",
                                                "unknown" };

    if (type <= AttributeRole::eUnknown)
    {
        return s << ognRoleNames[uint8_t(type)];
    }
    return s;
}

// Role names as used by USD, which are slightly different from the internal names used
inline std::string usdRoleName(const AttributeRole& type)
{
    static const std::string usdRoleNames[] = { "none",     "vector",        "normal",       "position",           "color",
                                                "texCoord", "quaternion",    "transform",    "frame",              "timecode",
                                                "text",     "appliedSchema", "primTypeName", "execution",          "matrix",
                                                "objectId", "bundle",        "path",         "instancedAttribute",  "ancestorPrimTypeName",
                                                "unknown" };


    if (type <= AttributeRole::eUnknown)
    {
        return usdRoleNames[uint8_t(type)];
    }
    return usdRoleNames[uint8_t(AttributeRole::eUnknown)];
}

struct Type
{
    BaseDataType baseType; // 1 byte

    // 1 for raw base types; 2 for vector2f, int2, etc; 3 for point3d, normal3f, etc;
    // 4 for quatf, color4d, float4, matrix2f etc; 9 for matrix3f, etc; 16 for matrix4d, etc.
    uint8_t componentCount; // 1 byte

    // 0 for a single value
    // 1 for an array
    // 2 for an array of arrays (not yet supported)
    uint8_t arrayDepth; // 1 byte

    AttributeRole role; // 1 byte

    constexpr Type(BaseDataType baseType, uint8_t componentCount = 1, uint8_t arrayDepth = 0, AttributeRole role = AttributeRole::eNone)
        : baseType(baseType), componentCount(componentCount), arrayDepth(arrayDepth), role(role)
    {
    }

    constexpr Type() : Type(BaseDataType::eUnknown)
    {
    }

    // Matches little endian interpretation of the four bytes
    constexpr explicit Type(const TypeC& t)
        : baseType(BaseDataType(t.type & 0xff)),
          componentCount((t.type >> 8) & 0xff),
          arrayDepth((t.type >> 16) & 0xff),
          role(AttributeRole((t.type >> 24) & 0xff))
    {
    }

    constexpr explicit operator TypeC() const
    {
        uint32_t type = uint8_t(role) << 24 | arrayDepth << 16 | componentCount << 8 | uint8_t(baseType);
        return TypeC{ type };
    }

    constexpr bool operator==(const Type& rhs) const
    {
        return compatibleRawData(rhs) && role == rhs.role;
    }

    constexpr bool operator!=(const Type& rhs) const
    {
        return !((*this) == rhs);
    }

    constexpr bool operator<(const Type& rhs) const
    {
        return TypeC(*this) < TypeC(rhs);
    }

    /**
     * Role-insensitive equality check
     */
    constexpr bool compatibleRawData(const Type& otherType) const
    {
        return baseType == otherType.baseType && componentCount == otherType.componentCount &&
               arrayDepth == otherType.arrayDepth;
    }

    /**
     * Check to see if this is one of the matrix types
     */
    constexpr bool isMatrixType() const
    {
        return (role == AttributeRole::eMatrix) || (role == AttributeRole::eFrame) || (role == AttributeRole::eTransform);
    }
    /**
     * Returns the dimensions of the type, componentCount for most types and square root of that for matrix types
     */
    constexpr uint8_t dimension() const
    {
        if (isMatrixType())
        {
            return componentCount == 4 ? 2 : (componentCount == 9 ? 3 : (componentCount == 16 ? 4 : componentCount));
        }
        return componentCount;
    }

    std::string getTypeName() const
    {
        std::ostringstream typeName;
        typeName << baseType;

        if (componentCount > 1)
            typeName << uint32_t(componentCount);

        if (arrayDepth == 1)
            typeName << "[]";
        else if (arrayDepth == 2)
            typeName << "[][]";

        // Some roles are hidden from USD
        if ((role != AttributeRole::eNone)
        &&  (role != AttributeRole::eObjectId)
        &&  (role != AttributeRole::eBundle)
        &&  (role != AttributeRole::ePath)
        )
        {
            typeName << " (" << usdRoleName(role) << ")";
        }

        return typeName.str();
    }
    // ======================================================================
    /**
     * OGN formats the type names slightly differently.
     *   - the tuples are internal "float[3]" instead of "float3"
     *   - the roles replace the actual name "colord[3]" instead of "double3 (color)"
     */
    std::string getOgnTypeName() const
    {
        std::ostringstream typeName;
        if (role == AttributeRole::eText)
        {
            typeName << "string";
            return typeName.str();
        }
        if (role == AttributeRole::ePath)
        {
            typeName << "path";
            return typeName.str();
        }
        if (role != AttributeRole::eNone)
        {
            typeName << role;
            // For roles with explicit types, add that to the role name
            if ((role != AttributeRole::eTimeCode)
            &&  (role != AttributeRole::eTransform)
            &&  (role != AttributeRole::eFrame)
            &&  (role != AttributeRole::eObjectId)
            &&  (role != AttributeRole::eBundle)
            &&  (role != AttributeRole::eExecution))
            {
                switch (baseType)
                {
                    case BaseDataType::eHalf:
                        typeName << "h";
                        break;
                    case BaseDataType::eFloat:
                        typeName << "f";
                        break;
                    case BaseDataType::eDouble:
                        typeName << "d";
                        break;
                    default:
                        typeName << baseType;
                        break;
                }
            }
        }
        else
        {
            typeName << baseType;
        }

        if (componentCount > 1)
        {
            typeName << "[" << uint32_t(dimension()) << "]";
        }

        if (arrayDepth == 1)
            typeName << "[]";
        else if (arrayDepth == 2)
            typeName << "[][]";

        return typeName.str();
    }
};

inline std::ostream& operator<<(std::ostream& s, const Type& type)
{
    s << type.getTypeName();
    return s;
}

}
}

namespace std
{
template <>
struct hash<carb::flatcache::Type>
{
    std::size_t operator()(const carb::flatcache::Type& key) const
    {
        return carb::flatcache::TypeC(key).type;
    }
};

template <>
struct hash<carb::flatcache::TypeC>
{
    std::size_t operator()(const carb::flatcache::TypeC& key) const
    {
        return key.type;
    }
};
}
