// Copyright (c) 2022, 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.
//
#include <OgnToUint64Database.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/graph/core/ogn/UsdTypes.h>

namespace omni
{
namespace graph
{
namespace nodes
{
namespace
{
template<typename T>
bool tryComputeAssumingType(OgnToUint64Database& db)
{
    auto functor = [](auto const& value, auto& converted) { converted = static_cast<uint64_t>(value); };
    return ogn::compute::tryComputeWithArrayBroadcasting<T, uint64_t>(db.inputs.value(), db.outputs.converted(), functor);
}

} // namespace

class OgnToUint64
{
public:
    // Node to convert numeric inputs to floats
    static bool compute(OgnToUint64Database& db)
    {
        auto& valueType = db.inputs.value().type();
        switch (valueType.baseType)
        {
            case BaseDataType::eBool:
                return tryComputeAssumingType<bool>(db);
            case BaseDataType::eUChar:
                return tryComputeAssumingType<uint8_t>(db);
            case BaseDataType::eInt:
                return tryComputeAssumingType<int32_t>(db);
            case BaseDataType::eUInt:
                return tryComputeAssumingType<uint32_t>(db);
            case BaseDataType::eInt64:
                return tryComputeAssumingType<int64_t>(db);
            case BaseDataType::eHalf:
                return tryComputeAssumingType<pxr::GfHalf>(db);
            case BaseDataType::eFloat:
                return tryComputeAssumingType<float>(db);
            case BaseDataType::eDouble:
                return tryComputeAssumingType<double>(db);
            default:
                db.logError("Failed to resolve input types");
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        auto valueAttr = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::value.token());
        auto outAttr = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::converted.token());

        auto valueType = valueAttr.iAttribute->getResolvedType(valueAttr);

        // The output shape must match the input shape and visa-versa, however we can't say anything
        // about the input base type until it's connected
        if (valueType.baseType != BaseDataType::eUnknown)
        {
            Type resultType(BaseDataType::eUInt64, 1, valueType.arrayDepth);
            outAttr.iAttribute->setResolvedType(outAttr, resultType);
        }
    }

};

REGISTER_OGN_NODE();

} // nodes
} // graph
} // omni
