// Copyright (c) 2022-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.
//
#include <OgnSelectIfDatabase.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <omni/graph/core/ogn/string.h>
#include <omni/graph/core/ogn/Types.h>
#include <carb/logging/Log.h>

namespace omni
{
namespace graph
{
namespace nodes
{

namespace
{
template<typename T>
bool tryComputeAssumingType(OgnSelectIfDatabase& db, size_t count)
{
    auto functor = [](auto const& a, auto const& b, auto const& condition, auto& result)
    {
        result = condition ? a : b;
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<T, T, bool, T>(
        db.inputs.ifTrue(), db.inputs.ifFalse(), db.inputs.condition(), db.outputs.result(), functor, count);
}

template<typename T, size_t N>
bool tryComputeAssumingType(OgnSelectIfDatabase& db, size_t count)
{
    auto functor = [](auto const& a, auto const& b, auto const& condition, auto& result)
    {
        if (condition)
        {
            memcpy(result, a, sizeof(T) * N);
        }
        else
        {
            memcpy(result, b, sizeof(T) * N);
        }
    };
    return ogn::compute::tryComputeWithArrayBroadcasting<T[N], T[N], bool, T[N]>(
        db.inputs.ifTrue(), db.inputs.ifFalse(), db.inputs.condition(), db.outputs.result(), functor, count);
}

} // namespace

class OgnSelectIf
{
public:
    static bool computeVectorized(OgnSelectIfDatabase& db, size_t count)
    {
        NodeObj nodeObj = db.abi_node();
        const AttributeObj ifTrueAttr = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::ifTrue.token());
        const Type ifTrueType = ifTrueAttr.iAttribute->getResolvedType(ifTrueAttr);
        const AttributeObj ifFalseAttr = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::ifFalse.token());
        const Type ifFalseType = ifFalseAttr.iAttribute->getResolvedType(ifFalseAttr);
        const AttributeObj conditionAttr = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::condition.token());
        const Type conditionType = conditionAttr.iAttribute->getResolvedType(conditionAttr);

        // This handles string case (arrays of strings are not supported)
        if (ifTrueType.baseType == BaseDataType::eUChar && ifTrueType.arrayDepth == 1
                && ifFalseType.baseType == BaseDataType::eUChar && ifFalseType.arrayDepth == 1
                && conditionType.arrayDepth == 0)
        {
            for (size_t idx = 0; idx < count; ++idx)
            {
                const bool condition = *db.inputs.condition(idx).template get<bool>();
                auto ifTrue = db.inputs.ifTrue(idx).template get<uint8_t[]>();
                auto ifFalse = db.inputs.ifFalse(idx).template get<uint8_t[]>();
                auto result = db.outputs.result(idx).template get<uint8_t[]>();
                if (condition)
                {
                    result->resize(ifTrue->size());
                    memcpy(&((*result)[0]), &((*ifTrue)[0]), result->size());
                }
                else
                {
                    result->resize(ifFalse->size());
                    memcpy(&((*result)[0]), &((*ifFalse)[0]), result->size());
                }
            }

            return true;
        }

        try
        {
            switch (ifTrueType.baseType)
            {
            case BaseDataType::eBool:
                return tryComputeAssumingType<bool>(db, count);
            case BaseDataType::eToken:
                return tryComputeAssumingType<ogn::Token>(db, count);
            case BaseDataType::eDouble:
                switch (ifTrueType.componentCount)
                {
                    case 1:  return tryComputeAssumingType<double>(db, count);
                    case 2:  return tryComputeAssumingType<double, 2>(db, count);
                    case 3:  return tryComputeAssumingType<double, 3>(db, count);
                    case 4:  return tryComputeAssumingType<double, 4>(db, count);
                    case 9:  return tryComputeAssumingType<double, 9>(db, count);
                    case 16: return tryComputeAssumingType<double, 16>(db, count);
                    default: break;
                }
            case BaseDataType::eFloat:
                switch (ifTrueType.componentCount)
                {
                    case 1: return tryComputeAssumingType<float>(db, count);
                    case 2: return tryComputeAssumingType<float, 2>(db, count);
                    case 3: return tryComputeAssumingType<float, 3>(db, count);
                    case 4: return tryComputeAssumingType<float, 4>(db, count);
                    default: break;
                }
            case BaseDataType::eHalf:
                switch (ifTrueType.componentCount)
                {
                    case 1: return tryComputeAssumingType<pxr::GfHalf>(db, count);
                    case 2: return tryComputeAssumingType<pxr::GfHalf, 2>(db, count);
                    case 3: return tryComputeAssumingType<pxr::GfHalf, 3>(db, count);
                    case 4: return tryComputeAssumingType<pxr::GfHalf, 4>(db, count);
                    default: break;
                }
            case BaseDataType::eInt:
                switch (ifTrueType.componentCount)
                {
                    case 1: return tryComputeAssumingType<int32_t>(db, count);
                    case 2: return tryComputeAssumingType<int32_t, 2>(db, count);
                    case 3: return tryComputeAssumingType<int32_t, 3>(db, count);
                    case 4: return tryComputeAssumingType<int32_t, 4>(db, count);
                    default: break;
                }
                ;
            case BaseDataType::eInt64:
                return tryComputeAssumingType<int64_t>(db, count);
            case BaseDataType::eUChar:
                return tryComputeAssumingType<unsigned char>(db, count);
            case BaseDataType::eUInt:
                return tryComputeAssumingType<uint32_t>(db, count);
            case BaseDataType::eUInt64:
                return tryComputeAssumingType<uint64_t>(db, count);
            default: break;
            }

            db.logWarning("Failed to resolve input types");
        }
        catch (ogn::compute::InputError &error)
        {
            db.logError(error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node){
        auto ifTrue = node.iNode->getAttributeByToken(node, inputs::ifTrue.token());
        auto ifFalse = node.iNode->getAttributeByToken(node, inputs::ifFalse.token());
        auto condition = node.iNode->getAttributeByToken(node, inputs::condition.token());
        auto result = node.iNode->getAttributeByToken(node, outputs::result.token());

        auto ifTrueType = ifTrue.iAttribute->getResolvedType(ifTrue);
        auto ifFalseType = ifFalse.iAttribute->getResolvedType(ifFalse);
        auto conditionType = condition.iAttribute->getResolvedType(condition);

        // Require ifTrue, ifFalse, and condition to be resolved before determining result's type
        if (ifTrueType.baseType != BaseDataType::eUnknown && ifFalseType.baseType != BaseDataType::eUnknown
            && conditionType.baseType != BaseDataType::eUnknown)
        {
            std::array<AttributeObj, 3> attrs { ifTrue, ifFalse, result };
            std::array<uint8_t, 3> tupleCounts {
                ifTrueType.componentCount,
                ifFalseType.componentCount,
                std::max(ifTrueType.componentCount, ifFalseType.componentCount)
            };
            std::array<uint8_t, 3> arrayDepths {
                ifTrueType.arrayDepth,
                ifFalseType.arrayDepth,
                std::max(conditionType.arrayDepth, std::max(ifTrueType.arrayDepth, ifFalseType.arrayDepth))
            };
            const bool isStringInput = (ifTrueType.baseType == BaseDataType::eUChar
                                    && ifTrueType.arrayDepth == 1
                                    && ifFalseType.baseType == BaseDataType::eUChar
                                    && ifFalseType.arrayDepth == 1
                                    && conditionType.arrayDepth == 0
                                    && ifTrueType.role == AttributeRole::eText
                                    && ifFalseType.role == AttributeRole::eText);
            std::array<AttributeRole, 3> rolesBuf {
                ifTrueType.role,
                ifFalseType.role,
                isStringInput ? AttributeRole::eText : AttributeRole::eUnknown
            };
            node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(),
                                                        arrayDepths.data(), rolesBuf.data(), attrs.size());
        }
    }
};

REGISTER_OGN_NODE()

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