// 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.
//
#pragma once
#include <carb/flatcache/IFlatcache.h>
#include <omni/graph/core/iComputeGraph.h>
#include <tuple>

namespace carb
{
namespace flatcache
{
using ArrayIndex = size_t;
}
}

namespace omni {
namespace graph {
namespace core {

// ====================================================================================================
// This is a temporary ABI interface used for prototyping the use of improved Gather. Some of these functions
// may move to other interfaces once the prototyping is complete. For now, none of them should be used outside
// of the prototyping efforts.
/*   _____          _   _       _     _    _
    |  __ \        | \ | |     | |   | |  | |
    | |  | | ___   |  \| | ___ | |_  | |  | |___  ___
    | |  | |/ _ \  | . ` |/ _ \| __| | |  | / __|/ _ \
    | |__| | (_) | | |\  | (_) | |_  | |__| \__ \  __/
    |_____/ \___/  |_| \_|\___/ \__|  \____/|___/\___|
*/

using PathBucketIndex = std::tuple<carb::flatcache::PathC, carb::flatcache::BucketId, carb::flatcache::ArrayIndex>;
using GatherId = uint64_t;
static constexpr GatherId kInvalidGatherId = 0;

/**
 * Gathered attributes can be automatically converted and copied to Hydra fast-path attributes
 */
enum class GatherAddTransformsMode
{
    eNone,  ///< do not add any attributes
    eLocal, ///< _localTransform
    eWorld  ///< _worldPosition, _worldOrientation, _worldScale
};

struct IGatherPrototype
{
    CARB_PLUGIN_INTERFACE("omni::graph::core::IGatherPrototype", 2, 0);

    /**
     * Adds the given paths to FC and updates the FC sync filter for this Graph. The prims may be tagged to ensure
     * the buckets do not contain unrelated prims. The returned GatherId identifies the FC buckets.
     *
     * After this call, the given paths will be present in FC, in a set of buckets which contain only these paths. This
     * is necessary to allow vectorized access to an attribute.
     *
     * @note The order of the paths supplied to this function may not be the order of the same paths in the Gather. Use
     *       getGatheredPaths() to get the gather-order of gathered prims.
     *
     * @note The returned GatherId can become invalidated if the underlying flatcache data changes, use isGatherValid()
     *       to check.
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] paths array of prim paths to add
     * @param[in] numPaths number of elements in paths
     * @param[in] allAttributes when true, all USD attributes will be gathered. When false, the "attributes" parameter
     *            is used.
     * @param[in] attributes array of attribute names to add to FC, when allAttributes is false
     * @param[in] numAttributes number of elements in attributes array
     * @param[in] addTransformsMode The transform attributes to create if any
     * @param[in] shouldWriteBack Flag to request to write back cached data to USD
     * @param[in] forceExportToHistory when true, all gathered paths will be tagged for being exported into the history
     *
     * @return The gather id corresponding to the gathered buckets, or kInvalidGatherId on failure.
     */
    GatherId(CARB_ABI* gatherPaths)(const GraphContextObj&,
                                    const carb::flatcache::PathC* paths,
                                    size_t numPaths,
                                    bool allAttributes,
                                    NameToken const* attributes,
                                    size_t numAttributes,
                                    GatherAddTransformsMode addTransformsMode,
                                    bool shouldWriteBack,
                                    bool forceExportToHistory);

    /**
     * Gets the array of paths in the given Gather, in gather order.
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] gatherId The Gather id returned from gatherPaths()
     * @param[out] refToPaths reference to the array of prim paths, in gather-order. This pointer is volatile, do not
     *             save. It can become invalid the next time something changes in flatcache.
     * @param[out] refToSize number of elements in refToPaths
     *
     * @return true if the Gather was found, false if it was not
     */
    bool (CARB_ABI*getGatheredPaths)(const GraphContextObj&, GatherId gatherId,
                                     carb::flatcache::PathC const*& refToPaths, size_t& refToSize);

    /**
     * Gets the array of buckets in the given Gather, in gather order.
     * For example if the gathered paths are P1,P2,P3,P4 which spans bucket B1, B2, we know that the sum of the
     * sizes of B1 and B2 must be 4. If B1 has 2 entries we know where the paths are located:
     *
     * B1
     * row 0: P1
     * row 1: P2
     *
     * B2
     * row 0: P3
     * row 1: P4
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] gatherId The Gather id returned from gatherPaths()
     * @param[out] refToBucketArray reference to the array of BucketIds, in gather-order. This pointer is volatile, do not
     *             save. It can become invalid the next time something changes in flatcache.
     * @param[out] refToSize number of elements in refToBucketArray
     *
     * @return true if the Gather was found, false if it was not
     */
    bool(CARB_ABI* getGatheredBuckets)(const GraphContextObj&,
                                     GatherId gatherId,
                                     BucketId const*& refToBucketArray,
                                     size_t& refToSize);

    /**
     * Gets the size information for the given gathered attribute name.
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] gatherId The Gather id returned from gatherPaths()
     * @param[in] attributeName name of the gathered attribute
     * @param[out] type The type information
     * @param[out] baseSizeBytes The size of the base elements of the Type
     * FIXME: Should be accessible from flatcache ABI?
     * @return true if the attribute was found, false if it was not
     */
    bool(CARB_ABI* getGatheredType)(const GraphContextObj&, GatherId gatherId, NameToken attributeName, Type& type, size_t& baseSizeBytes);

    // ==============================================================================================================
    // The functions below here are temporary replacements for the same functions removed from IGraphContext. They
    // were being modified to use attributes, causing a breaking change, and since they were going to be changed anyway
    // they were moved to this prototype for easier update later.
    // ==============================================================================================================

    /**
     * Given an attribute connected to a gather node, retrieves the attribute value in system memory
     * Deprecated: to be replaced with the 2.0 API soon
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The attribute structure to which the operation applies
     * @param[in] flags The data access flags indicating whether the data is to be R, W, or RW
     * @return the void pointer to the data
     */
    void*(CARB_ABI* getGatherArray)(const GraphContextObj& context, const AttributeObj& attrObj, DataAccessFlags flags);

    /**
     * Given an attribute connected to a gather node, retrieves the attribute value in GPU memory
     * Deprecated: to be replaced with the 2.0 API soon
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The attribute structure to which the operation applies
     * @param[in] flags The data access flags indicating whether the data is to be R, W, or RW
     * @return the void pointer to the data
     */
    void*(CARB_ABI* getGatherArrayGPU)(const GraphContextObj& context, const AttributeObj& attrObj, DataAccessFlags flags);

    /**
     * Given an attribute connected to a gather node, retrieves the array of the paths it has gathered
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The source attribute structure. This can be any of the gather node's attributes.
     * @return the pointer to the path array, encoded as uint64_t. These integers can be cast to SdfPaths.
     */
    const uint64_t*(CARB_ABI* getGatherPathArray)(const GraphContextObj& context, const AttributeObj& attrObj);

    /**
     * Given an array attribute in a gather node, retrieves an array of the array sizes in CPU memory
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The source attribute structure.
     * @param[in] flags The data access flags indicating whether the data is to be R, W, or RW
     * @return the pointer to the array
     */
    size_t*(CARB_ABI* getGatherArrayAttributeSizes)(const GraphContextObj& context,
                                                    const AttributeObj& attrObj,
                                                    DataAccessFlags flags);

    /**
     * Given an array attribute in a gather node, retrieves an array of the array sizes in GPU memory
     * Deprecated: to be replaced with the 2.0 API soon
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The source attribute structure.
     * @param[in] flags The data access flags indicating whether the data is to be R, W, or RW
     * @return the pointer to the array
     */
    size_t*(CARB_ABI* getGatherArrayAttributeSizesGPU)(const GraphContextObj& context,
                                                       const AttributeObj& attrObj,
                                                       DataAccessFlags flags);

    /**
     * Given a gather node attribute, retrieves the number of elements
     * Deprecated: to be replaced with the 2.0 API soon
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] attrObj The attribute structure to which the operation applies
     * @return the number of elements gathered, -1 if operation is unsuccessful
     */
    size_t(CARB_ABI* getElementCount)(const GraphContextObj& context, const AttributeObj& attrObj);

    /**
     * Gets the array of repeated paths and where to find its data in the given Gather.
     * Each element in the array will have the path of the prim, the bucket id of it's data, and the index of its position inside the bucket
     *
     * @param[in] context structure containing both the interface and underlying object
     * @param[in] gatherId The Gather id returned from gatherPaths()
     * @param[out] refToRepeatedPathsArray reference to the array of PathBucketIndex which outlines the location of the
     *             data for repeated paths in the gater. This pointer is volatile, do not
     *             save. It can become invalid the next time something changes in flatcache.
     * @param[out] refToSize number of elements in refToRepeatedPathsArray
     *
     * @return true if the Gather was found, false if it was not
     */
    bool(CARB_ABI* getGatheredRepeatedPaths)(const GraphContextObj&,
                                             GatherId gatherId,
                                             PathBucketIndex const*& refToRepeatedPathsArray,
                                             size_t& refToSize);

};
// Update this every time a new ABI function is added, to ensure one isn't accidentally added in the middle
STRUCT_INTEGRITY_CHECK(IGatherPrototype, getGatheredRepeatedPaths, 10)

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