// Copyright (c) 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.
//
//

#include <OgnGatherByPathDatabase.h>

#include <omni/graph/core/IGatherPrototype.h>
#include <carb/flatcache/IPath.h>
#include <algorithm>
#include <vector>

using namespace carb::flatcache;


namespace omni
{
namespace graph
{
namespace core
{

class OgnGatherByPath
{
    // The arguments from the last compute call - we want to skip our compute if these haven't changed, and we
    // haven't heard that our Gather has changed.
    struct LastComputeState
    {
        // copies of input values
        std::vector<NameToken> primPathTokens; // Used in compute to compare to input tokens
        bool allAttributes{ false };
        std::string attributes;
        NameToken hydraFastPath{ 0 };
        bool shouldWriteBack;
        bool forceExportToHistory{ false };

        // data use by path-changed callback
        std::vector<PathC> primPaths; // Used in OG callback to compare with given paths
        PathChangedCallback callback{ nullptr, nullptr };
        bool checkResyncAttributes{ false }; // Set to true to be notified for resync children
        bool gatherDirty{ false }; // Set to true by OG callback

        const carb::flatcache::IPath * iPath = nullptr;
    };

public:
    LastComputeState m_state;

    // Called by OG when anything changes - which indicates buckets may have been modified
    static void onPathChanged(const carb::flatcache::PathC* paths, const size_t numPaths, void* userData)
    {
        NodeHandle nodeHandle = reinterpret_cast<NodeHandle>(userData);
        auto& state =
            OgnGatherByPathDatabase::sm_stateManagerOgnGatherByPath.getState<OgnGatherByPath>(nodeHandle).m_state;
        // set the dirty flag if we recognize one of the changed prim paths - which will indicate a 'resync' notice
        // we don't care about attribute values changing
        const auto& gatheredPaths = state.primPaths;
        auto iter = std::find_first_of(gatheredPaths.begin(), gatheredPaths.end(), paths, paths + numPaths);
        state.gatherDirty = (iter != gatheredPaths.end());
        if (!state.gatherDirty && state.checkResyncAttributes && state.iPath)
        {
            auto& iPath = state.iPath;
            auto iter = std::find_first_of(gatheredPaths.begin(), gatheredPaths.end(), paths, paths + numPaths,
                                           [iPath](const PathC a, const PathC b) { return a == iPath->getParent(b); });
            state.gatherDirty = (iter != gatheredPaths.end());
        }
    }

    // ----------------------------------------------------------------------------
    static void initialize(const GraphContextObj& context, const NodeObj& nodeObj)
    {
        auto& cppObj = OgnGatherByPathDatabase::sInternalState<OgnGatherByPath>(nodeObj);
        cppObj.m_state.callback = { &onPathChanged, reinterpret_cast<void*>(nodeObj.nodeHandle) };
        nodeObj.iNode->registerPathChangedCallback(nodeObj, cppObj.m_state.callback);
    }

    // ----------------------------------------------------------------------------
    static void release(const NodeObj& nodeObj)
    {
        auto& cppObj = OgnGatherByPathDatabase::sInternalState<OgnGatherByPath>(nodeObj);
        nodeObj.iNode->deregisterPathChangedCallback(nodeObj, cppObj.m_state.callback);
    }

    // ----------------------------------------------------------------------------
    static bool compute(OgnGatherByPathDatabase& db)
    {
        const omni::graph::core::IGatherPrototype* iGatherPrototype =
                carb::getCachedInterface<omni::graph::core::IGatherPrototype>();

        const IGraphContext& iContext = *db.abi_context().iContext;

        auto& cppObj = db.internalState<OgnGatherByPath>();
        LastComputeState& state = cppObj.m_state;

        const auto& primPathTokens = db.inputs.primPaths();
        bool allAttributes = db.inputs.allAttributes();
        const auto& attributes = db.inputs.attributes();
        auto hydraFastPathToken = db.inputs.hydraFastPath();
        const bool& shouldWriteBack = db.inputs.shouldWriteBack();
        bool forceExportToHistory = db.inputs.forceExportToHistory();

        if (primPathTokens.empty())
            return true;

        state.checkResyncAttributes = db.inputs.checkResyncAttributes();
        state.iPath = db.abi_context().iPath;

        // We can skip this compute if our state isn't dirty AND none of our inputs have changed
        if (!state.gatherDirty
            && (allAttributes == state.allAttributes)
            && (attributes == state.attributes)
            && (hydraFastPathToken == state.hydraFastPath)
            && (primPathTokens.size() == state.primPathTokens.size())
            && (shouldWriteBack == state.shouldWriteBack)
            && (forceExportToHistory == state.forceExportToHistory)
            && std::equal(primPathTokens.begin(), primPathTokens.end(), state.primPathTokens.begin()))
        {
            return true;
        }

        // convert the given NameToken[] to PathC[]
        std::vector<PathC> paths;
        paths.resize(primPathTokens.size());
        std::transform(primPathTokens.begin(), primPathTokens.end(), paths.begin(),
                       [&](const auto& p) { return db.abi_context().iPath->getHandle(db.tokenToString(p)); });

        // Save the input args to our per-instance state
        state.primPaths.resize(paths.size());
        state.primPathTokens.resize(paths.size());

        std::copy(paths.begin(), paths.end(), state.primPaths.begin());
        std::copy(primPathTokens.begin(), primPathTokens.end(), state.primPathTokens.begin());

        state.allAttributes = allAttributes;
        state.attributes = attributes;
        state.hydraFastPath = hydraFastPathToken;
        state.shouldWriteBack = shouldWriteBack;
        state.forceExportToHistory = forceExportToHistory;

        // tokenize the given list of space-separated attribute names
        std::vector<NameToken> gatherAttribs;
        if (!attributes.empty())
        {
            if (allAttributes)
            {
                CARB_LOG_ERROR("'attributes' specified, but allAttributes is true");
                return false;
            }
            std::istringstream iss(attributes);
            std::string word;
            while (std::getline(iss, word, ' '))
            {
                if (!word.empty())
                    gatherAttribs.push_back(db.stringToToken(word.c_str()));
            }
        }

        GatherAddTransformsMode hydraTransformsMode = GatherAddTransformsMode::eNone;
        if (hydraFastPathToken == db.tokens.Local)
            hydraTransformsMode = GatherAddTransformsMode::eLocal;
        else if (hydraFastPathToken == db.tokens.World)
            hydraTransformsMode = GatherAddTransformsMode::eWorld;

        // Make the Gather ABI call
        GatherId gatherId = iGatherPrototype->gatherPaths(
            db.abi_context(), paths.data(), paths.size(),
            allAttributes, gatherAttribs.data(), gatherAttribs.size(), hydraTransformsMode, shouldWriteBack, forceExportToHistory);

        db.outputs.gatherId() = gatherId;

        // If the gather succeeded, copy the gathered paths into the output
        if (gatherId != kInvalidGatherId)
        {
            size_t gatheredPathsSize{ 0 };
            PathC const* gatheredPaths{ nullptr };
            size_t repeatedPathsSize{ 0 };
            PathBucketIndex const* repeatedPaths{ nullptr };

            if (iGatherPrototype->getGatheredPaths(db.abi_context(), gatherId, gatheredPaths, gatheredPathsSize) &&
                iGatherPrototype->getGatheredRepeatedPaths(db.abi_context(), gatherId, repeatedPaths, repeatedPathsSize))
            {
                db.outputs.gatheredPaths().resize(gatheredPathsSize + repeatedPathsSize);
                std::transform(gatheredPaths, gatheredPaths + gatheredPathsSize, db.outputs.gatheredPaths().begin(),
                               [&db](PathC const& path) {
                                   Path pathWrap(path);
                                   return db.stringToToken(pathWrap.getText());
                               });
                std::transform(repeatedPaths, repeatedPaths + repeatedPathsSize, db.outputs.gatheredPaths().begin() + gatheredPathsSize,
                               [&db](PathBucketIndex const& pathBucketIndex) {
                                   Path pathWrap(std::get<0>(pathBucketIndex));
                                   return db.stringToToken(pathWrap.getText());
                               });
                // clear the state gather dirty flag
                state.gatherDirty = false;
                return true;
            }
        }

        db.outputs.gatheredPaths().resize(0);
        CARB_LOG_ERROR("Could not gather paths");

        return false;
    }
};

REGISTER_OGN_NODE()

} // namespace core
} // namespace graph
} // namespace omni
