
// 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 "FindPrimsPathTracker.h"
#include "LayerIdentifierResolver.h"
#include "OgnWritePrimsV2Database.h"
#include "SpanUtils.h"
#include "WritePrimCommon.h"

namespace omni
{
namespace graph
{
namespace nodes
{

class OgnWritePrimsV2
{
    // keep a set of prims that has been *created* by this node.
    // If the node writes to an existing prim on the stage, it's not counted as created.
    // it will be reset every time removeMissingPrims is true to track the diff
    // track Usd and Fabric separately since usdWriteBack can be turned on or off
    PXR_NS::SdfPathUSet createdUsdPrims;
    PXR_NS::SdfPathUSet createdFabricPrims;

    // keep a map of attributes that has been *created* by this node.
    // If the node writes to an existing attribute on the stage, it's not counted as created.
    // it will be reset every time removeMissingAttributes is true to track the diff
    // track Usd and Fabric separately since usdWriteBack can be turned on or off
    PXR_NS::SdfPathUSet createdUsdAttributes;
    PXR_NS::SdfPathUSet createdFabricAttributes;

    PathToMetadataMap createdObjectMetadata;

    PrimPathsState targetPrimPaths;

    static size_t getConnectedPrimCount(OgnWritePrimsV2Database& db)
    {
        GraphContextObj const& contextObj = db.abi_context();
        NodeObj const& nodeObj = db.abi_node();

        NodeContextHandle const nodeHandle = nodeObj.nodeContextHandle;
        NameToken const primsToken = inputs::primsBundle.token();
        return contextObj.iContext->getInputTargetCount(contextObj, nodeHandle, primsToken, db.getInstanceIndex());
    }

    //! Check if any of the targetPrim paths are descendants of any of the Root Prims paths in find/read prims.
    //! If any of the path is an descendant, then the output prim(s) can potentially mutate the input prims' states,
    //! creating a loop. This is undesirable and prohibited to prevent user errors.
    //!
    //! returns true if loop is detected.
    static bool detectFindPrimsPathsLoop(GraphContextObj const& context,
                                         const NodeObj& nodeObj,
                                         PathVector const& targetPrimPaths)
    {
        fabric::FabricId fabricId;
        std::tie(fabricId, std::ignore) = getFabricForNode(context, nodeObj);

        auto const rootPrimPaths = getRootPrimPathsSet(fabricId);
        if (rootPrimPaths.empty())
        {
            return false;
        }

        // AbsoluteRootPath has to be the smallest path in the set if exists
        // a quick test
        if (*rootPrimPaths.begin() == PXR_NS::SdfPath::AbsoluteRootPath())
        {
            ogn::OmniGraphDatabase::logError(
                nodeObj,
                "FindPrims/ReadPrims source contains absolute root path \"/\", writing to any target on the stage may mutate the input state. "
                "Please choose a different source root prim path for FindPrims/ReadPrims.");
            return true;
        }

        PXR_NS::SdfPathSet targetPrimPathsSet;
        std::transform(targetPrimPaths.cbegin(), targetPrimPaths.cend(),
                       std::inserter(targetPrimPathsSet, targetPrimPathsSet.begin()),
                       [](fabric::PathC pathC) { return fabric::toSdfPath(pathC); });

        for (auto const& rootPrimPath : rootPrimPaths)
        {
            auto const lb = targetPrimPathsSet.lower_bound(rootPrimPath);
            // if a lower_bound in targetPrimPaths is a descendant of rootPrimPath, a loop could form
            if (lb != targetPrimPathsSet.end() && lb->HasPrefix(rootPrimPath))
            {
                ogn::OmniGraphDatabase::logError(
                    nodeObj,
                    "Target prim '%s' is a descendant of FindPrims/ReadPrims source '%s', writing to the target may mutate the input state. Please choose a different target",
                    lb->GetText(), rootPrimPath.GetText());
                return true;
            }
        }

        return false;
    }

    static void handleMissingPrims(OgnWritePrimsV2Database& db,
                                   PXR_NS::SdfPathUSet const& createdPrimPaths,
                                   bool removeMissingPrims = true)
    {
        auto& internalState = OgnWritePrimsV2Database::sInternalState<OgnWritePrimsV2>(db.abi_node());
        auto& prevCreatedUsdPrims = internalState.createdUsdPrims;
        auto& prevCreatedFabricPrims = internalState.createdFabricPrims;
        // if removeMissingPrims and usdWriteBack, find the diff of this compute's createdPrimPaths vs last
        // removeMissingPrims's createdPrimPaths, and removes the ones that no longer exist in current
        // createdPrimPaths
        auto processMissingPrims = [&db, &createdPrimPaths, removeMissingPrims](
                                       PXR_NS::SdfPathUSet& prevCreatedPrims, bool isUsd)
        {
            if (createdPrimPaths == prevCreatedPrims)
            {
                return;
            }

            if (removeMissingPrims)
            {
                // if removeMissingPrims is on:
                // - always removes missing prims from fabric
                // - if usdWriteBack is on, also remove missing prims from USD

                PXR_NS::SdfPathVector primPathsToRemove;
                primPathsToRemove.reserve(prevCreatedPrims.size());
                for (const auto& path : prevCreatedPrims)
                {
                    // previously created prim no longer exists
                    if (createdPrimPaths.find(path) == createdPrimPaths.end())
                    {
                        ogn::OmniGraphDatabase::logWarning(
                            db.abi_node(),
                            "Previously created prim %s no longer exists in upstream bundle. Removing it from %s...",
                            path.GetText(), isUsd ? "USD" : "Fabric");

                        primPathsToRemove.push_back(path);
                    }
                }

                if (!primPathsToRemove.empty())
                {
                    removePrims(db.abi_context(), db.abi_node(), primPathsToRemove, isUsd);
                }

                // now the accumulated created paths should be createdPrimPaths from this execution
                prevCreatedPrims = createdPrimPaths;
            }
            else
            {
                // else accumulate the written paths until next removeMissingPrims happens
                prevCreatedPrims.insert(createdPrimPaths.begin(), createdPrimPaths.end());
            }
        };

        // prims are always written to Fabric
        processMissingPrims(prevCreatedFabricPrims, false);

        // only when usdWriteBack is true will prims be written to USD
        if (db.inputs.usdWriteBack())
        {
            processMissingPrims(prevCreatedUsdPrims, true);
        }
    }

    static void handleMissingAttributes(OgnWritePrimsV2Database& db,
                                        PXR_NS::SdfPathUSet const& createdAttributes,
                                        bool removeMissingAttrs = true)
    {
        auto& internalState = OgnWritePrimsV2Database::sInternalState<OgnWritePrimsV2>(db.abi_node());
        auto& prevCreatedUsdAttributes = internalState.createdUsdAttributes;
        auto& prevCreatedFabricAttributes = internalState.createdFabricAttributes;

        auto processMissingAttributes = [&db, &createdAttributes, removeMissingAttrs](
                                            PXR_NS::SdfPathUSet& prevCreatedAttributes, bool isUsd)
        {
            if (createdAttributes == prevCreatedAttributes)
            {
                return;
            }

            if (removeMissingAttrs)
            {
                // if removeMissingAttrs is on:
                // - always removes missing attributes from fabric
                // - if usdWriteBack is on, also remove missing attributes from USD

                PXR_NS::SdfPathVector attrPathsToRemove;
                attrPathsToRemove.reserve(prevCreatedAttributes.size());
                for (const auto& path : prevCreatedAttributes)
                {
                    // previously created prim no longer exists
                    if (createdAttributes.find(path) == createdAttributes.end())
                    {
                        ogn::OmniGraphDatabase::logWarning(
                            db.abi_node(),
                            "Previously created attribute %s no longer exists in upstream bundle. Removing it from %s...",
                            path.GetText(), isUsd ? "USD" : "Fabric");

                        attrPathsToRemove.push_back(path);
                    }
                }

                if (!attrPathsToRemove.empty())
                {
                    removeAttributes(db.abi_context(), db.abi_node(), attrPathsToRemove, isUsd);
                }

                // now the accumulated created paths should be createdAttributes from this execution
                prevCreatedAttributes = createdAttributes;

            }
            else
            {
                // else accumulate the written attributes until next removeMissingAttrs happens
                prevCreatedAttributes.insert(createdAttributes.begin(), createdAttributes.end());
            }
        };

        // attributes are always written to Fabric
        processMissingAttributes(prevCreatedFabricAttributes, false);

        // only when usdWriteBack is true will attributes be written to USD
        if (db.inputs.usdWriteBack())
        {
            processMissingAttributes(prevCreatedUsdAttributes, true);
        }
    }

    static void handleMissingMetadata(OgnWritePrimsV2Database& db,
                                      PathToMetadataMap const& createdObjectMetadata,
                                      bool removeMissingMetadata = true)
    {
        auto& internalState = OgnWritePrimsV2Database::sInternalState<OgnWritePrimsV2>(db.abi_node());
        auto& prevCreatedObjectMetadata = internalState.createdObjectMetadata;

        auto processMissingMetadata = [&db, &createdObjectMetadata, removeMissingMetadata](
                                       PathToMetadataMap& prevCreatedObjectMetadata, bool isUsd)
        {
            if (createdObjectMetadata == prevCreatedObjectMetadata)
            {
                return;
            }

            if (removeMissingMetadata)
            {
                PathToMetadataMap metadataToRemove;

                for (const auto& objectMetadata : prevCreatedObjectMetadata)
                {
                    const auto& objectPath = objectMetadata.first;
                    const auto& prevMetadataKeys = objectMetadata.second;

                    if (createdObjectMetadata.find(objectPath) == createdObjectMetadata.end())
                    {
                        ogn::OmniGraphDatabase::logWarning(
                            db.abi_node(),
                            "Previously created metadata for object %s no longer exists in upstream bundle. Removing it from %s...",
                            objectPath.GetText(), isUsd ? "USD" : "Fabric");

                        metadataToRemove[objectPath] = prevMetadataKeys;
                    }
                    else
                    {
                        const auto& metadataKeys = createdObjectMetadata.at(objectPath);

                        for (const auto& key : prevMetadataKeys)
                        {
                            if (metadataKeys.find(key) == metadataKeys.end())
                            {
                                ogn::OmniGraphDatabase::logWarning(
                                    db.abi_node(),
                                    "Previously created metadata %s for object %s no longer exists in upstream bundle. Removing it from %s...",
                                    key.GetText(), objectPath.GetText(), isUsd ? "USD" : "Fabric");

                                metadataToRemove[objectPath].insert(key);
                            }
                        }
                    }
                }
                if (!metadataToRemove.empty())
                {
                    removeMetadata(db.abi_context(), db.abi_node(), metadataToRemove, isUsd);
                }

                prevCreatedObjectMetadata = createdObjectMetadata;
            }
            else
            {
                for (const auto& objectMetadata : createdObjectMetadata)
                {
                    // accumulate created metadata until this function is invoked with removeMissingAttrs
                    prevCreatedObjectMetadata[objectMetadata.first].insert(objectMetadata.second.begin(), objectMetadata.second.end());
                }
            }
        };

        // XXX: metadata export for fabric not yet implemented
        //processMissingMetadata(prevCreatedObjectMetadata, false);

        if (db.inputs.usdWriteBack())
        {
            processMissingMetadata(prevCreatedObjectMetadata, true);
        }
    }

    static WritePrimResult scatterUnderTargets(GraphContextObj const& contextObj,
                                               NodeObj const& nodeObj,
                                               gsl::span<ConstBundleHandle> connectedBundleHandles,
                                               PathVector const& targetPrimPaths,
                                               bool const usdWriteBack,
                                               NameToken const& layerIdentifier,
                                               WritePrimMatchers const& matchers,
                                               WrittenPrimAndAttributesTracker& primAndAttrsTracker,
                                               WritePrimDiagnostics& diagnostics)
    {
        auto result = WritePrimResult::None;

        // a map to contain all valid bundle prims info from ALL connected bundles and descendant bundles.
        BundlePrimInfoMap bundlePrimInfoMap;

        for (ConstBundleHandle const& bundleHandle : connectedBundleHandles)
        {
            WritePrimResult const subResult = collectPrimsFromBundleHierarchy(
                contextObj, nodeObj, diagnostics, bundleHandle, matchers, bundlePrimInfoMap);

            mergeWritePrimResult(result, subResult);
        }

        auto writeUnderTarget = [&contextObj, &nodeObj, &diagnostics, usdWriteBack, layerIdentifier, &result,
                                 &bundlePrimInfoMap, &primAndAttrsTracker](PathC const& targetPrefix)
        {
            for (auto const& [_, info] : bundlePrimInfoMap)
            {
                WritePrimResult const subResult = writePrimAndAttributes(
                    contextObj, nodeObj, diagnostics, info, usdWriteBack, layerIdentifier, targetPrefix,
                    fabric::asInt(PXR_NS::SdfPath::AbsoluteRootPath()), primAndAttrsTracker);

                mergeWritePrimResult(result, subResult);
            }
        };

        for (auto const& target : targetPrimPaths)
        {
            writeUnderTarget(target);
        }

        return result;
    }

    static WritePrimResult writeToTargets(GraphContextObj const& contextObj,
                                          NodeObj const& nodeObj,
                                          gsl::span<ConstBundleHandle> connectedBundleHandles,
                                          PathVector const& targetPrimPaths,
                                          bool const usdWriteBack,
                                          NameToken const& layerIdentifier,
                                          WritePrimMatchers const& matchers,
                                          WrittenPrimAndAttributesTracker& primAndAttrsTracker,
                                          WritePrimDiagnostics& diagnostics)
    {
        auto result = WritePrimResult::None;
        auto& internalState = OgnWritePrimsV2Database::sInternalState<OgnWritePrimsV2>(nodeObj);

        size_t targetPrimIndex = 0;
        for (ConstBundleHandle const& bundleHandle : connectedBundleHandles)
        {
            // a map to contain all valid bundle prims info from ONE connected bundle and descendant bundles.
            BundlePrimInfoMap bundlePrimInfoMap;

            WritePrimResult const subResult = collectPrimsFromBundleHierarchy(
                contextObj, nodeObj, diagnostics, bundleHandle, matchers, bundlePrimInfoMap);

            mergeWritePrimResult(result, subResult);

            // remove the paths that are already prefixed by other paths.
            // similar to SdfPath::RemoveDescendentPaths, but BundlePrimInfoMap is already ordered so we don't have to
            // sort again
            PXR_NS::SdfPathVector uniqueRootPaths;
            uniqueRootPaths.reserve(bundlePrimInfoMap.size());
            for (auto const& entry : bundlePrimInfoMap)
            {
                if (uniqueRootPaths.empty() || !entry.first.HasPrefix(uniqueRootPaths.back()))
                {
                    uniqueRootPaths.push_back(entry.first);
                }
            }

            // index to the currently mapped unique root prim path
            size_t rootIndex = 0;
            for (auto const& [sourcePrimPath, info] : bundlePrimInfoMap)
            {
                // when the sourcePrimPath is no longer a descendant of current unique root, move it onto the next one
                if (rootIndex < uniqueRootPaths.size() && !sourcePrimPath.HasPrefix(uniqueRootPaths[rootIndex]))
                {
                    rootIndex++;

                    // each root index takes up a target prim index
                    targetPrimIndex++;
                }

                if (targetPrimIndex >= targetPrimPaths.size())
                {
                    ogn::OmniGraphDatabase::logWarning(
                        nodeObj,
                        "Not enough 'prims' targets to write out all source prims. Skipping the rest of the prims");
                    break;
                }

                WritePrimResult const subResult =
                    writePrimAndAttributes(contextObj, nodeObj, diagnostics, info, usdWriteBack, layerIdentifier,
                                           targetPrimPaths[(PathVector::size_type)targetPrimIndex],
                                           fabric::asInt(uniqueRootPaths[rootIndex]), primAndAttrsTracker);

                mergeWritePrimResult(result, subResult);
            }

            // always increment targetPrimIndex when moving onto next bundle
            targetPrimIndex++;
        }

        return result;
    }

public:
    static bool compute(OgnWritePrimsV2Database& db)
    {
        auto result = WritePrimResult::None;

        auto& context = db.abi_context();
        auto& nodeObj = db.abi_node();

        auto& internalState = OgnWritePrimsV2Database::sInternalState<OgnWritePrimsV2>(nodeObj);

        bool inputChanged = db.state.attrNamesToExport() != db.inputs.attrNamesToExport() ||
                            db.state.pathPattern() != db.inputs.pathPattern() ||
                            db.state.typePattern() != db.inputs.typePattern() ||
                            db.state.usdWriteBack() != db.inputs.usdWriteBack() ||
                            db.state.scatterUnderTargets() != db.inputs.scatterUnderTargets() ||
                            db.state.layerIdentifier() != db.inputs.layerIdentifier();

        db.state.attrNamesToExport() = db.inputs.attrNamesToExport();
        db.state.pathPattern() = db.inputs.pathPattern();
        db.state.typePattern() = db.inputs.typePattern();
        db.state.usdWriteBack() = db.inputs.usdWriteBack();
        db.state.scatterUnderTargets() = db.inputs.scatterUnderTargets();
        db.state.layerIdentifier() = db.inputs.layerIdentifier();

        PathVector const targetPrimPaths =
            getPathsFromInputRelationship(context, nodeObj, inputs::prims.token(), db.getInstanceIndex());
        if (internalState.targetPrimPaths.changed(targetPrimPaths))
        {
            internalState.targetPrimPaths.store(targetPrimPaths);
            inputChanged = true;
        }

#if 0
        // TODO dirty tracking of primBundle

        if (!inputChanged)
        {
            return true;
        }
#endif

        long stageId = context.iContext->getStageId(context);
        auto stage = PXR_NS::UsdUtilsStageCache::Get().Find(PXR_NS::UsdStageCache::Id::FromLongInt(stageId));

        auto resolvedLayerIdentifier =
            resolveLayerIdentifier(nodeObj, stage, inputs::layerIdentifier.m_token, db.inputs.layerIdentifier());
        
        // Since the bundle input allows multiple connections, we need to collect all of them
        size_t const connectedPrimCount = getConnectedPrimCount(db);

        PXR_NS::SdfPathUSet createdPrimPaths;
        PXR_NS::SdfPathUSet createdAttributes;
        PathToMetadataMap createdObjectMetadata;

        // Even when we have no connected prims, we still return true, because the node *did* compute, and execOut
        // attribute will be set.
        if (connectedPrimCount > 0)
        {
            auto handler = [&db, &result, &targetPrimPaths, &createdPrimPaths, &createdAttributes, &createdObjectMetadata,
                            &internalState, &resolvedLayerIdentifier](gsl::span<ConstBundleHandle> connectedBundleHandles)
            {
                GraphContextObj const& contextObj = db.abi_context();
                NodeObj const& nodeObj = db.abi_node();

                if (detectFindPrimsPathsLoop(contextObj, nodeObj, targetPrimPaths))
                {
                    mergeWritePrimResult(result, WritePrimResult::Fail);

                    // do not proceed
                    return;
                }

                NodeContextHandle const nodeHandle = nodeObj.nodeContextHandle;
                NameToken const primsToken = inputs::primsBundle.token();

                bool const usdWriteBack = db.inputs.usdWriteBack();
                ogn::const_string const attrPattern = db.inputs.attrNamesToExport();
                ogn::const_string const pathPattern = db.inputs.pathPattern();
                ogn::const_string const typePattern = db.inputs.typePattern();

                WritePrimMatchers const matchers{ attrPattern, pathPattern, typePattern };

                contextObj.iContext->getInputTargets(contextObj, nodeHandle, primsToken,
                                                     (omni::fabric::PathC*)(connectedBundleHandles.data()),
                                                     db.getInstanceIndex());

                // Write back all bundles
                WritePrimDiagnostics diagnostics;
                // NOTE: writing metadata is currently only implemented for USD writeback,
                // so no distinction in made for USD vs Fabric metadata collections.
                WrittenPrimAndAttributesTracker tracker(
                    createdPrimPaths, usdWriteBack ? internalState.createdUsdPrims : internalState.createdFabricPrims,
                    createdAttributes,
                    usdWriteBack ? internalState.createdUsdAttributes : internalState.createdFabricAttributes,
                    createdObjectMetadata,
                    internalState.createdObjectMetadata);

                auto writeResult = WritePrimResult::None;
                if (db.inputs.scatterUnderTargets())
                {
                    writeResult =
                        scatterUnderTargets(contextObj, nodeObj, connectedBundleHandles, targetPrimPaths, usdWriteBack,
                                            resolvedLayerIdentifier, matchers, tracker, diagnostics);
                }
                else
                {
                    writeResult = writeToTargets(contextObj, nodeObj, connectedBundleHandles, targetPrimPaths,
                                                 usdWriteBack, resolvedLayerIdentifier, matchers, tracker, diagnostics);
                }

                // Report skipped invalid bundles
                mergeWritePrimResult(result, diagnostics.report(contextObj, nodeObj));
            };

            withStackScopeBuffer<ConstBundleHandle>(connectedPrimCount, handler);
        }

        {
            PXR_NS::SdfChangeBlock changeblock;
            // Track primBundles content changes and remove no longer existing attributes or prims 
            // TODO when dirtyId checking works correctly, only handle them when upstream bundle or related inputs
            // changed.
            handleMissingPrims(db, createdPrimPaths);
            // TODO ReadPrims node will *not* remove an attribute from the bundle if it's removed due to lack of USD
            // change tracking. The only way to trigger removal of missing attributes currently is to change filter
            // pattern or disconnect then reconnect primBundles connection to force a re-evaluation.
            handleMissingAttributes(db, createdAttributes);

            handleMissingMetadata(db, createdObjectMetadata);
        }

        auto const ok = result != WritePrimResult::Fail;
        if (ok)
        {
            db.outputs.execOut() = kExecutionAttributeStateEnabled;
        }

        return ok;
    }
};

REGISTER_OGN_NODE()

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