// 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.
//

// clang-format off
#include "UsdPCH.h"
// clang-format on

#include "PrimCommon.h"
#include "VariantCommon.h"

#include <carb/logging/Log.h>

#include <omni/fabric/Enums.h>
#include <omni/graph/core/ogn/Types.h>
#include <omni/math/linalg/SafeCast.h>
#include <omni/usd/UsdContext.h>
#include <pxr/usd/sdf/variantSetSpec.h>
#include <pxr/usd/sdf/variantSpec.h>

#include <OgnBlendVariantsDatabase.h>

using omni::graph::core::ogn::eAttributeType::kOgnOutput;
using omni::graph::core::ogn::eMemoryType::kCpu;
using omni::math::linalg::vec3f;

using DB = OgnBlendVariantsDatabase;

namespace omni::graph::nodes
{
namespace
{
using LerpFunction = void (*)(const double alpha,
                              const pxr::SdfPropertySpecHandle& propertyA,
                              const pxr::SdfPropertySpecHandle& propertyB,
                              pxr::UsdAttribute attribute);
template <typename T>
void floatLerp(const double alpha,
          const pxr::SdfPropertySpecHandle& propertyA,
          const pxr::SdfPropertySpecHandle& propertyB,
          pxr::UsdAttribute attribute)
{
    T a = propertyA->GetDefaultValue().Get<T>();
    T b = propertyB->GetDefaultValue().Get<T>();
    T c = pxr::GfLerp(alpha, a, b);
    if (attribute.IsValid())
        attribute.Set<T>(c);
}

template <typename T>
void discreteLerp(const double alpha,
            const pxr::SdfPropertySpecHandle& propertyA,
            const pxr::SdfPropertySpecHandle& propertyB,
            pxr::UsdAttribute attribute)
{
    T a = propertyA->GetDefaultValue().Get<T>();
    T b = propertyB->GetDefaultValue().Get<T>();
    if (attribute.IsValid())
    {
        if (alpha < 0.5)
            attribute.Set<T>(a);
        else
            attribute.Set<T>(b);
    }
}

void lerpAttribute(const double alpha,
                   const pxr::SdfPropertySpecHandle& propertyA,
                   const pxr::SdfPropertySpecHandle& propertyB,
                   const pxr::UsdAttribute& attribute)
{
    if (propertyA->GetSpecType() != propertyB->GetSpecType())
        throw warning("Property spec types do not match (attribute " + attribute.GetPath().GetString() + ")");

    if (propertyA->GetTypeName() != attribute.GetTypeName())
        throw warning("Attribute types do not match (attribute " + attribute.GetPath().GetString() + ")");

    if (propertyA->GetValueType() != propertyB->GetValueType())
        throw warning("Property value types do not match");

    auto typeName = propertyA->GetValueType().GetTypeName();

    auto handleType = [alpha, &propertyA, &propertyB, &attribute, &typeName](const char* type, LerpFunction lerpFunction) -> bool
    {
        if (typeName == type)
        {
            lerpFunction(alpha, propertyA, propertyB, attribute);
            return true;
        }
        return false;
    };


    if (!handleType("bool", discreteLerp<bool>)
        &&  !handleType("double", floatLerp<double>)
        &&  !handleType("float", floatLerp<float>)
        &&  !handleType("pxr_half::half", floatLerp<pxr::GfHalf>)
        &&  !handleType("int", discreteLerp<int>)
        &&  !handleType("__int64", discreteLerp<int64_t>) // Windows
        &&  !handleType("long", discreteLerp<int64_t>) // Linux
        &&  !handleType("unsigned char", discreteLerp<uint8_t>)
        &&  !handleType("unsigned int", discreteLerp<uint32_t>)
        &&  !handleType("unsigned __int64", discreteLerp<uint64_t>) // Windows
        &&  !handleType("unsigned long", discreteLerp<uint64_t>) // Linux
        &&  !handleType("TfToken", discreteLerp<pxr::TfToken>)
        &&  !handleType("SdfTimeCode", floatLerp<pxr::SdfTimeCode>)
        &&  !handleType("GfVec2d", floatLerp<pxr::GfVec2d>)
        &&  !handleType("GfVec2f", floatLerp<pxr::GfVec2f>)
        &&  !handleType("GfVec2h", floatLerp<pxr::GfVec2h>)
        &&  !handleType("GfVec2i", discreteLerp<pxr::GfVec2i>)
        &&  !handleType("GfVec3d", floatLerp<pxr::GfVec3d>)
        &&  !handleType("GfVec3f", floatLerp<pxr::GfVec3f>)
        &&  !handleType("GfVec3h", floatLerp<pxr::GfVec3h>)
        &&  !handleType("GfVec3i", discreteLerp<pxr::GfVec3i>)
        &&  !handleType("GfVec4d", floatLerp<pxr::GfVec4d>)
        &&  !handleType("GfVec4f", floatLerp<pxr::GfVec4f>)
        &&  !handleType("GfVec4h", floatLerp<pxr::GfVec4h>)
        &&  !handleType("GfVec4i", discreteLerp<pxr::GfVec4i>)
        &&  !handleType("GfQuatd", floatLerp<pxr::GfQuatd>)
        &&  !handleType("GfQuatf", floatLerp<pxr::GfQuatf>)
        &&  !handleType("GfQuath", floatLerp<pxr::GfQuath>)
        &&  !handleType("GfMatrix2d", floatLerp<pxr::GfMatrix2d>)
        &&  !handleType("GfMatrix3d", floatLerp<pxr::GfMatrix3d>)
        &&  !handleType("GfMatrix4d", floatLerp<pxr::GfMatrix4d>))
        throw warning("Unsupported property type " + typeName);
}
}

class OgnBlendVariants
{
public:
    static bool compute(OgnBlendVariantsDatabase& db)
    {
        auto ok = [&db]()
        {
            db.outputs.execOut() = kExecutionAttributeStateEnabled;
            return true;
        };

        try
        {
            pxr::UsdPrim prim = tryGetTargetPrim(db, db.inputs.prim(), "prim");

            std::string variantSetName = db.tokenToString(db.inputs.variantSetName());
            std::string variantNameA = db.tokenToString(db.inputs.variantNameA());
            std::string variantNameB = db.tokenToString(db.inputs.variantNameB());
            double blend = std::max(std::min(db.inputs.blend(), 1.0), 0.0);

            pxr::UsdVariantSets variantSets = prim.GetVariantSets();
            pxr::UsdVariantSet variantSet = variantSets.GetVariantSet(variantSetName);
            if (!variantSet.IsValid())
                throw warning("Invalid variant set " + variantSetName);

            bool finishing = (1.0 - blend) < 1e-6;

            if (finishing && db.inputs.setVariant())
                variantSet.SetVariantSelection(variantNameB);

            VariantData a = getVariantData(prim, variantSetName, variantNameA);
            VariantData b = getVariantData(prim, variantSetName, variantNameB);

            for (const auto& [path, propertyA] : a)
            {
                if (b.find(path) == b.end())
                    continue;

                auto propertyB = b[path];

                auto attribute = prim.GetStage()->GetAttributeAtPath(path);
                if (!attribute.IsValid())
                    throw warning("Invalid attribute " + path.GetString());

                if (finishing && db.inputs.setVariant())
                    attribute.Clear();
                else
                    lerpAttribute(blend, propertyA, propertyB, attribute);
            }
            return ok();
        }
        catch (const warning& e)
        {
            db.logWarning(e.what());
        }
        catch (const std::exception& e)
        {
            db.logError(e.what());
        }

        return false;
    }
};

REGISTER_OGN_NODE()

} // namespace omni::graph::nodes
