// Copyright (c) 2021-2021, 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 <OgnReadPrimAttributeDatabase.h>

#include <omni/graph/core/IAttributeType.h>
#include <omni/graph/core/StringUtils.h>
#include <omni/graph/core/CppWrappers.h>
#include <omni/usd/UsdContext.h>
#include <omni/fabric/FabricUSD.h>

#include "PrimCommon.h"
#include "ReadPrimCommon.h"

using namespace omni::fabric;

namespace omni
{
namespace graph
{
namespace nodes
{
static thread_local bool recurseGuard{ false };

class OgnReadPrimAttribute
{
    static bool importAtTime(omni::fabric::FabricId fabricId,
                      NodeObj const& nodeObj,
                      GraphContextObj const& context,
                      pxr::UsdPrim const& prim,
                      TokenC srcName,
                      pxr::SdfTimeCode const& time,
                      OgnReadPrimAttributeDatabase& db,
                      size_t offset)
    {
        CARB_PROFILE_ZONE(1, "ImportAtTime");
        IFabricUsd* iFabricUsd = carb::getCachedInterface<IFabricUsd>();

        // Handle the case where the prim is not yet available
        if (! prim.IsValid())
        {
            CARB_LOG_WARN("Attempted to import time from invalid prim from node '%s' and source '%s'",
                nodeObj.iNode->getPrimPath(nodeObj), omni::fabric::toTfToken(srcName).GetText()
            );
            return false;
        }

        pxr::SdfPath importPath = prim.GetPath();
        bool force = false;
        if (UsdTimeCode::Default() != time)
        {
            force = true;
            char const* primPath = nodeObj.iNode->getPrimPath(nodeObj);
            static constexpr char const kImportSubPath[] = "/Import";
            std::string storingPath;
            storingPath.reserve(strlen(primPath) + sizeof(kImportSubPath));
            storingPath = primPath;
            storingPath += kImportSubPath;
            importPath = pxr::SdfPath(storingPath);
        }

        {
            CARB_PROFILE_ZONE(1, "ImportUSDToFabric");
            BucketId srcPrimBucketId =
                omni::graph::nodes::addAttribToCache(prim, importPath, srcName, iFabricUsd, fabricId, time, force);

            if (srcPrimBucketId == kInvalidBucketId)
            {
                OgnReadPrimAttributeDatabase::logError(
                    nodeObj, "Unable to add prim %s to fabric", prim.GetPath().GetText());
                return false;
            }
        }

        db.state.time(offset) = time.GetValue();
        db.state.srcPath(offset) = asInt(prim.GetPath()).path;
        db.state.srcPathAsToken(offset) = asInt(prim.GetPath().GetAsToken()).token;
        db.state.importPath(offset) = asInt(importPath).path;
        db.state.srcAttrib(offset) = srcName.token;

        return true;
    }

    static void setup(NodeObj const& nodeObj,
               GraphContextObj const& context,
               pxr::SdfTimeCode const& time,
               OgnReadPrimAttributeDatabase& db,
               size_t offset,
               bool isInCompute)
    {
        if (recurseGuard)
            return;

        CARB_PROFILE_ZONE(1, "setup");
        recurseGuard = true;
        std::shared_ptr<nullptr_t> atScopExist(nullptr, [](auto) { recurseGuard = false; });

        auto typeInterface{ carb::getCachedInterface<omni::graph::core::IAttributeType>() };
        InstanceIndex instanceIndex = db.getInstanceIndex() + offset;

        db.state.correctlySetup(offset) = false;

        pxr::UsdAttribute attrib;
        try
        {
            attrib = omni::graph::nodes::findSelectedAttribute(context, nodeObj, instanceIndex);
        }
        catch (std::runtime_error const& error)
        {
            if (isInCompute)
                OgnReadPrimAttributeDatabase::logError(nodeObj,error.what());
            return;
        }
        if (!attrib)
        {
            // Couldn't get the indicated attribute for some expected reason
            if (isInCompute)
                OgnReadPrimAttributeDatabase::logError(nodeObj, "Prim has not been set");
            return;
        }

        auto typeName{ attrib ? attrib.GetTypeName().GetAsToken() : pxr::TfToken() };
        Type attribType{ typeInterface->typeFromSdfTypeName(typeName.GetText()) };
        tryResolveOutputAttribute(nodeObj, outputs::value.m_token, attribType);

        // Get interfaces
        const INode& iNode = *nodeObj.iNode;

        pxr::UsdPrim prim{ attrib.GetPrim() };

        if (!attrib.IsAuthored())
        {

            if (UsdTimeCode::Default() != time)
            {
                CARB_LOG_ERROR("Access a non authored attrib %s:%s with a specific timestamp is not allowed",
                               attrib.GetName().GetText(), prim.GetPath().GetText());
                return;
            }

            // fabric only cares about _authored_ attributes, so in order to make use of fabric mirrored arrays we
            // have to contrive to have values authored, even if we just want to read them.
            //
            // FIXME: OM-36961 exists to work around this issue some other way.
            CARB_LOG_INFO("%s: Attribute %s is not authored, authoring now", iNode.getPrimPath(nodeObj),
                          attrib.GetPath().GetText());

            // --------------------------------------------------------------------------------------
            // Get PrimVar into fabric by creating it on the prim
            if (pxr::UsdGeomPrimvar::IsPrimvar(attrib))
            {
                try
                {
                    addPrimvarToFabric(prim, attrib);
                }
                catch (std::runtime_error const& error)
                {
                    OgnReadPrimAttributeDatabase::logError(nodeObj,error.what());
                    return;
                }
            }
            else
            {
                // --------------------------------------------------------------------------------------
                // Get non-primvar into fabric by authoring the attribute with the composed value
                pxr::VtValue val;
                if (!attrib.Get(&val))
                {
                    OgnReadPrimAttributeDatabase::logError(nodeObj,"Unable to read %s of type %s", attrib.GetPath().GetText(), val.GetTypeName().c_str());
                    return;
                }

                if (!attrib.Set(val))
                {
                    OgnReadPrimAttributeDatabase::logError(nodeObj,"Unable to author %s with value of type %s", attrib.GetPath().GetText(),
                                val.GetTypeName().c_str());
                    return;
                }
            }
        }

        // Fill Fabric with USD data
        pxr::UsdStageRefPtr stage;
        omni::fabric::FabricId fabricId;
        std::tie(fabricId, stage) = getFabricForNode(context, nodeObj);
        if (fabricId == omni::fabric::kInvalidFabricId)
            return;

        if (!importAtTime(fabricId, nodeObj, context, prim, asInt(attrib.GetName()), time, db, offset))
            return;

        // If it's resolved, we already know that it is compatible from the above check of the USD
        AttributeObj outAttrib = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::value.m_token);
        Type outType = outAttrib.iAttribute->getResolvedType(outAttrib);
        if (outType == Type())
        {
            // Not resolved, so we have to resolve it now. This node is strange in that the resolved output type
            // depends on external state instead of other attributes.
            outAttrib.iAttribute->setResolvedType(outAttrib, attribType);
        }

        db.state.correctlySetup(offset) = true;
    }

    /**
     * Adds the given primvar attrib to fabric.At the moment this requires explicit authoring of the value
     * @param prim: The prim which is in FC, but doesn't have the given primvar
     * @param attrib: The primvar attribute for the prim, which has no value
     * @throws: std::runtime_error on failure
     * @
     */
    static void addPrimvarToFabric(pxr::UsdPrim prim, pxr::UsdAttribute& attrib)
    {
        // Primvars have to go through prim var api
        if (auto primVarAPI = pxr::UsdGeomPrimvarsAPI(prim))
        {
            if (pxr::UsdGeomPrimvar primVar = primVarAPI.FindPrimvarWithInheritance(attrib.GetName()))
            {
                if (auto newPrimVar = primVarAPI.CreatePrimvar(attrib.GetName(), primVar.GetTypeName()))
                {
                    if (auto& newAttr = newPrimVar.GetAttr())
                    {
                        // Get the inherited primvar attrib value
                        pxr::VtValue val;
                        if (primVar.Get(&val))
                        {
                            if (!newAttr.Set(val))
                            {
                                throw std::runtime_error(formatString("Unable to author primvar %s with value of type %s",
                                            newAttr.GetPath().GetText(), val.GetTypeName().c_str()));
                            }
                        }
                        else
                        {
                            // We don't have any value to copy, so lets use the default
                            auto const& defaultValue = primVar.GetTypeName().GetDefaultValue();
                            if (!newAttr.Set(defaultValue))
                            {
                                throw std::runtime_error(
                                    formatString("Unable to author primvar %s with default value", newAttr.GetPath().GetText()));
                            }
                        }
                    }
                }
            }
        }
    }

    // ----------------------------------------------------------------------------
    // Called by OG when our prim attrib changes. We want to catch the case of changing the prim attribute interactively
    static void onValueChanged(AttributeObj const& attrObj, void const* userData)
    {
        // FIXME: Be pedantic about validity checks - this can be run directly by the TfNotice so who knows
        // when or where this is happening
        NodeObj nodeObj{ attrObj.iAttribute->getNode(attrObj) };
        if (nodeObj.nodeHandle == kInvalidNodeHandle)
            return;

        GraphObj graphObj{ nodeObj.iNode->getGraph(nodeObj) };
        if (graphObj.graphHandle == kInvalidGraphHandle)
            return;

        //no setup if graph is disabled (during python edition for instance)
        if (graphObj.iGraph->isDisabled(graphObj))
            return;

        GraphContextObj context{ graphObj.iGraph->getDefaultGraphContext(graphObj) };
        if (context.contextHandle == kInvalidGraphContextHandle)
            return;

        OgnReadPrimAttributeDatabase db(context, nodeObj);
        setup(nodeObj, context, db.state.time(), db, 0, false);
    }

public:

    // ----------------------------------------------------------------------------
    static void initialize(GraphContextObj const& context, NodeObj const& nodeObj)
    {
        // We need to check resolution if any of our relevant inputs change
        std::array<NameToken, 4> attribNames{ inputs::name.m_token, inputs::usePath.m_token, inputs::primPath.m_token,
                                              inputs::prim.m_token };
        for (auto const& attribName : attribNames)
        {
            AttributeObj attribObj = nodeObj.iNode->getAttributeByToken(nodeObj, attribName);
            attribObj.iAttribute->registerValueChangedCallback(attribObj, onValueChanged, true);
        }
    }

    // ----------------------------------------------------------------------------
    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        GraphObj graphObj{ nodeObj.iNode->getGraph(nodeObj) };
        if (graphObj.graphHandle == kInvalidGraphHandle)
            return;

        GraphContextObj context{ graphObj.iGraph->getDefaultGraphContext(graphObj) };
        if (context.contextHandle == kInvalidGraphContextHandle)
            return;

        OgnReadPrimAttributeDatabase db(context, nodeObj);
        setup(nodeObj, context, db.state.time(), db, 0, false);
    }

    // ----------------------------------------------------------------------------
    static bool computeVectorized(OgnReadPrimAttributeDatabase& db, size_t count)
    {
        NodeObj nodeObj = db.abi_node();
        GraphContextObj ctx = db.abi_context();

        pxr::UsdStageRefPtr stage;
        omni::fabric::FabricId fabricId;
        std::tie(fabricId, stage) = getFabricForNode(ctx, nodeObj);
        if (fabricId == omni::fabric::kInvalidFabricId)
            return false;

        //////////////////////////////////////////////////////////////////////////
        //Gather vectorized data
        auto time = db.inputs.usdTimecode.vectorized(count);
        auto usePath = db.inputs.usePath.vectorized(count);
        auto primPath = db.inputs.primPath.vectorized(count);
        auto name = db.inputs.name.vectorized(count);

        auto correctlySetup = db.state.correctlySetup.vectorized(count);
        auto stateTime = db.state.time.vectorized(count);
        auto importPath = db.state.importPath.vectorized(count);
        auto srcPath = db.state.srcPath.vectorized(count);
        auto srcPathAsToken = db.state.srcPathAsToken.vectorized(count);
        auto srcAttrib = db.state.srcAttrib.vectorized(count);


        //////////////////////////////////////////////////////////////////////////
        // Make sure setup is correctly performed
        for (size_t idx = 0; idx < count; ++idx)
        {
            if (!correctlySetup[idx])
            {
                setup(nodeObj, ctx, time[idx], db, idx, true);
            }
            else
            {
                auto pathToken = usePath[idx] ? primPath[idx] : db.pathToToken(db.inputs.prim.firstOrDefault(idx));;
                if (pathToken != srcPathAsToken[idx] || name[idx] != srcAttrib[idx])
                    setup(nodeObj, ctx, time[idx], db, idx, true);
            }

            if (correctlySetup[idx] && UsdTimeCode(time[idx]) != stateTime[idx])
            {
                auto prim = stage->GetPrimAtPath(intToPath(srcPath[idx]));
                correctlySetup[idx] = importAtTime(fabricId, nodeObj, ctx, prim, name[idx], time[idx], db, idx);
            }
        }

        //////////////////////////////////////////////////////////////////////////
        //Do the copy by the data model or by hand if possible

        // for array or GPU attributes, go through the data model to properly handle the copy
        Type const type = db.outputs.value().type();
        bool const useDataModel = type.arrayDepth || ctx.iAttributeData->gpuValid(db.outputs.value().abi_handle(), ctx);

        if (useDataModel)
        {
            for (size_t idx = 0; idx < count; ++idx)
            {
                if (correctlySetup[idx])
                {
                    ConstAttributeDataHandle const src{ AttrKey{ importPath[idx], srcAttrib[idx] } };
                    db.abi_context().iAttributeData->copyData(
                        db.outputs.value(idx).abi_handle(), ctx, src);
                }
            }
        }
        else
        {
            // gather dst base pointer
            uint8_t* dstData = nullptr;
            size_t stride = 0;
            db.outputs.value().rawData(dstData, stride);

            // retrieve src pointers
            std::vector<ConstAttributeDataHandle> srcHandles;
            srcHandles.resize(count);
            for (size_t idx = 0; idx < count; ++idx)
            {
                srcHandles[idx] = correctlySetup[idx] ?
                                      ConstAttributeDataHandle(AttrKey{ importPath[idx], srcAttrib[idx] }) :
                                      ConstAttributeDataHandle::invalidHandle();
            }

            std::vector<void const*> srcData;
            srcData.resize(count);
            {
                CARB_PROFILE_ZONE(1, "Retrieve Pointers");
                ctx.iAttributeData->getDataR(srcData.data(), ctx, srcHandles.data(), count);
            }

            for (size_t idx = 0; idx < count; ++idx)
            {
                if (correctlySetup[idx])
                {
                    void* dst = dstData + idx * stride;
                    if (srcData[idx])
                        memcpy(dst, srcData[idx], stride);
                    else
                        db.logError("computeVectorized attempting to memcpy from nullptr");
                }
            }
        }


        return true;
    }

    // ----------------------------------------------------------------------------
    static bool updateNodeVersion(const GraphContextObj& context, const NodeObj& nodeObj, int oldVersion, int newVersion)
    {
        if (oldVersion < newVersion)
        {
            if (oldVersion < 2)
            {
                return upgradeUsdTimeCodeInput(context, nodeObj);
            }
        }
        return false;
    }
};

REGISTER_OGN_NODE()
} // namespace nodes
} // namespace graph
} // namespace omni
