// Copyright (c) 2019-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.
//
#pragma once

#include <carb/Interface.h>
#include <carb/flatcache/PathToAttributesMap.h>

namespace carb
{
namespace flatcache
{

// Callers of createCache() and getCache() can store anything they want in
// UserId. For example, OmniGraph uses it to store the OmniGraph pointer.
struct UserId
{
    uint64_t id;

    bool operator<(UserId b) const
    {
        return id < b.id;
    }
    bool operator==(UserId b) const
    {
        return id == b.id;
    }
    bool operator!=(UserId b) const
    {
        return id != b.id;
    }
};

constexpr UserId kDefaultUserId = { 0 };
constexpr UserId kInvalidUserId = { ~uint64_t(0) };

// Flatcache has the option to save a finite number of frames of history,
// organized as a ringbuffer. This is typically used to buffer data between
// simulation rendering. The simplest case, double buffering, allows simulation
// and rendering to run in parallel, each running for the full frame.
// Longer buffers can be used to feed one or more renderers running at
// different rates to simulation.
// To enable this feature, pass CacheType::eWithHistory to createCache(),
// otherwise pass CacheType::eWithoutHistory.
// Multiple caches can be created for each UsdStageId, but at most one can have
// history.
enum class CacheType
{
    eWithHistory,
    eWithoutHistory,
    eWithoutHistoryAndWithCuda,
    eWithHistoryAndCuda
};

struct FlatCache
{
    CARB_PLUGIN_INTERFACE("carb::flatcache::FlatCache", 0, 4);

    // Abstractly, a flatcache maps USD paths to USD attributes, just like a
    // UsdStage does.
    // Concretely we represent a flatcache by objects of type PathToAttributesMap.
    // This method creates a PathToAttributesMap for a given stage, but doesn't
    // populate it with values. This allows the cache to be filled lazily as
    // values are needed.
    // Instead, it traverses the given Usd stage making an index of paths to
    // attributes.
    // The cache uses the index to organize data into contiguous arrays,
    // and also allows you to find prims by type and/or attribute without
    // traversing the stage.
    // This method also specifies the stage to be used by calls to usdToCache()
    // and cacheToUsd().

    PathToAttributesMap&(CARB_ABI* createCache)(UsdStageId usdStageId, UserId userId, CacheType cacheType);


    void(CARB_ABI* addPrimToCache)(PathToAttributesMap& cache, const pxr::UsdPrim& prim, const std::set<TokenC>& filter);

    // Destroy the cache associated with the given stage.
    void(CARB_ABI* destroyCache)(UsdStageId usdStageId, UserId userId);

    // Prefetch the whole USD stage to the cache
    // Typically you only call this at stage load time, because the USD notify
    // handler updates the cache if the stage changes.
    void(CARB_ABI* usdToCache)(PathToAttributesMap& cache);

    // Write back all dirty cached data to the USD stage.
    // If your renderer doesn't use the cache then you need to do this
    // before rendering.
    void(CARB_ABI* cacheToUsd)(PathToAttributesMap& cache);

    // Write back only one bucket to usd
    void(CARB_ABI* cacheBucketToUsd)(PathToAttributesMap& cache, BucketId bucketId, bool skipMeshPoints);

    TypeC(CARB_ABI* usdTypeToTypeC)(pxr::SdfValueTypeName usdType);

    PathToAttributesMap*(CARB_ABI* getCache)(UsdStageId usdStageId, UserId userId);

    pxr::SdfValueTypeName(CARB_ABI* typeCtoUsdType)(TypeC typeC);

    size_t(CARB_ABI* getUsdTypeCount)();

    void(CARB_ABI* getAllUsdTypes)(TypeC* outArray, size_t outArraySize);

    /** @brief Import a prim in cache
     */
    void(CARB_ABI* addPrimToCacheNoOverwrite)(PathToAttributesMap& cache,
                                              const pxr::UsdPrim& prim,
                                              const std::set<TokenC>& filter);

    void(CARB_ABI* initStaticVariables)();

    void(CARB_ABI* exportUsd)(PathToAttributesMap& cache, pxr::UsdStageRefPtr usdStage, const double* timeCode, const double* prevTimeCode);

    /** @brief Attempt to serialize the cache into the specified buffer.
    *
    * @cache[in] The cache to serialize
    * @dest[in/out] Pointer to buffer to be written to, will start writing to head
    *       of pointer. dest will be left pointing to the point after the last write
    * @destSize Size of buffer that was allocated for the data (in bytes)
    * @pathStringCache - looking up strings is slow, yo
    *
    * @return Number of bytes written success is determined by (return <= @destSize)
    *
    *
    * @invariant It is safe to write to any memory within[dest, dest+size] for the
    *            duration of the function call.
    *
    * @note If the cache will not fit into the size of memory allocated in
    *       @dest then it will stop writing, but continue to run the serialize
    *       algorithm to calculate the actual amount of data that needs to be
    *       written
    *
    * @Todo : make cache const - not possible because serializeMirroredAray is not
    *         const, however, that is because getArraySpanC is used which also doesn't
    *         have a const version, so that needs to be addressed first, this is because
    *         in the call stack we end up with a copy from GPU -> CPU which would need to
    *         be avoided
    */
    uint64_t(CARB_ABI* serializeCache)(PathToAttributesMap& cache, uint8_t* dest, size_t destSize, SerializationCache& pathStringCache);

    /** @brief Given a buffer that has the serialized version of a cache written
     *         using the serialize function, this function will override all the data
     *         in the cache with the data from the buffer
     *
     * @cache[in/out] Reference to the cache to be populated
     * @pathCache[in/out] Looking up SDFPath via string can be expensive to it
     *                    is worthwhile to cache this data across many repeated
     *                    calls.
     * @input[in] Pointer to buffer of data containing serialized cache
     * @inputSize[in] Size of data in the buffer
     * @skipStageConfirmation[in] Whether we should skip making sure the destination stage is open.
     *
     * @return True if buffer was successfully de-serialized
     *
     * @note : this currently has to clear the cache before it is populated which is a possibly
     *         expensive operation
     *
     * @TODO: whould we care that it came from the same version of the USD file?
     */
    bool(CARB_ABI* deserializeCache)(
        PathToAttributesMap& destStage,
        DeserializationCache& pathCache,
        const uint8_t* input,
        const size_t inputSize,
        bool skipStageConfirmation);

    /** @brief Write a cache file to disk at a specified location
     *
     * @note many parameters to this function are optional

     * @cache[in] That cache to be written to disk
     * @file[in The location the file is desired to be written to
     * @workingBuffer[in] [Optional] In order to avoid costly reallocations
     *                    the code will attempt to re-use the memory at the buffer
     *                    location if it is large enough. If the buffer isn't larg
     *                    enough the cost of allocation, and re-traversal may be paid
     * @workingBufferSize[in] [Optional] If workingBuffer is non null, then this desrcibes the length
     *                        of the buffer
     * @return The amount of data needed to serialize the cache, a return value of 0 indicates an error
     *
     */
    uint64_t(CARB_ABI* writeCacheToDisk)(PathToAttributesMap& cache,
        const char* file,
        uint8_t* workingBuffer,
        uint64_t workingBufferSize);

    /** @brief Read a cache file from the specified location
     *
     * @file[in] The location the file is desired to be written to
     * @cache[in/out] That cache to be populates
     * @pathCache[in/out] Looking up SDFPath via string can be expensive to it
     *                    is worthwhile to cache this data across many repeated
     *                    calls.
     * @buffer[in/out] Buffer to use to read the cache file in to, passed to
     *                 allow reuse than allocate per call. Will be resized if not large enough.

     * @return Whether the read was successful
     *
     */
    bool(CARB_ABI* readCacheFromDisk)(PathToAttributesMap& cache,
        const char* fileName,
        DeserializationCache& pathCache,
        std::vector<uint8_t>& buffer);

    /** @brief Enable/Disable change notifications on USD changes.
     *
     * @enable[in] True/False enable notifications
     *
     */
    void(CARB_ABI* setEnableChangeNotifies)(bool enable);

    /** @brief Return whether change notifications on USD changes is enabled.
     *
     * @return True if change notifications on USD changes is enabled, else False.
     *
     */
    bool(CARB_ABI* getEnableChangeNotifies)();

    /** @brief make buckets for all prims on a USD stage, but only if this
    *          hasn't been done before.
    *
    * This is used to lazily create an index of all prims on a stage, without
    * the time or memory cost of fetching all the attribute values. The user
    * can then use findPrims to, for example, find all the prims of a
    * particular type.
    * 
    * If a SimStageWithHistory hasn't been created for this stage then a
    * warning will be printed and no population will be done.
    *
    * @cache[in] The PathToAttributesMap to populate
    */
    void(CARB_ABI* minimalPopulateIfNecessary)(PathToAttributesMap& cache);

};
}
}
