// 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.
//

#pragma once

#include "ConstBundlePrims.h"

#include <omni/graph/core/CppWrappers.h>
#include <omni/graph/core/iComputeGraph.h>
#include <omni/graph/core/ComputeGraph.h>

#include <algorithm>

namespace omni
{
namespace graph
{
namespace io
{

// ====================================================================================================
//
//                                      Const Bundle Primitive
//
// ====================================================================================================

inline ConstBundlePrim::ConstBundlePrim(ConstBundlePrims& bundlePrims,
                                        BundlePrimIndex primIndex,
                                        omni::core::ObjectPtr<omni::graph::core::IConstBundle2> bundle)
    : m_bundlePrims{ &bundlePrims }, m_bundle{ std::move(bundle) }, m_primIndex{ primIndex }
{
    // Read and cache all non internal attributes.
    readAndCacheAttributes();
}

inline void ConstBundlePrim::readAndCacheAttributes() noexcept
{
    using namespace omni::graph::core;

    IConstBundle2* bundle = getConstBundlePtr();
    GraphContextObj const& context = getConstBundlePrims()->context();

    std::vector<ConstAttributeDataHandle> attrHandles(bundle->getAttributeCount());
    bundle->getConstAttributes(attrHandles.data(), attrHandles.size());

    auto& attrs = getAttributes();
    for(ConstAttributeDataHandle& attrHandle : attrHandles)
    {
        if(!attrHandle.isValid())
            continue;

        NameToken attrName = context.iAttributeData->getName(context, attrHandle);
        attrs.insert(std::make_pair(attrName, std::make_unique<BundleAttrib>(*this, attrName)));
    }
}

inline BundleAttrib const* ConstBundlePrim::getConstAttr(core::NameToken attrName) noexcept
{
    using namespace omni::graph::core;

    // Try to find cached attributes
    auto& attrMap = getAttributes();
    auto it = attrMap.find(attrName);
    if (it != attrMap.end())
    {
        return it->second.get();
    }

    // Try to find attribute in this bundle.
    IConstBundle2* bundle = getConstBundlePtr();
    ConstAttributeDataHandle attributeHandle = bundle->getConstAttributeByName(attrName);
    if (!attributeHandle.isValid())
    {
        // attribute is not found, ensure entry is removed from the cache.
        auto it = attrMap.find(attrName);
        if (it != attrMap.end())
        {
            attrMap.erase(it);
        }

        return nullptr;
    }

    // Check if attribute in the bundle is stale
    auto newPrimAttribute = new BundleAttrib{ *this, attrName};
    std::unique_ptr<BundleAttrib> primAttributePtr{ newPrimAttribute };
    attrMap.emplace(attrName, std::move(primAttributePtr));
    return newPrimAttribute;
}

inline BundleAttrib const* ConstBundlePrim::getAttr(omni::graph::core::NameToken attrName) const noexcept
{
    return const_cast<ConstBundlePrim*>(this)->getConstAttr(attrName);
}

inline size_t ConstBundlePrim::attrCount() noexcept
{
    return getAttributes().size();
}

inline BundlePrimIndex ConstBundlePrim::primIndex() noexcept
{
    return m_primIndex;
}

inline omni::graph::core::NameToken ConstBundlePrim::path() noexcept
{
    using namespace omni::graph::core;

    ConstBundlePrims* bundlePrims = getConstBundlePrims();
    BundlePrimIndex const primIndex = this->primIndex();
    if (primIndex >=  bundlePrims->getPrimCount())
    {
        return carb::flatcache::kUninitializedToken;
    }
    NameToken const* paths = bundlePrims->getConstPrimPaths();
    return (paths != nullptr) ? paths[primIndex] : carb::flatcache::kUninitializedToken;
}

inline omni::graph::core::NameToken ConstBundlePrim::path() const noexcept
{
    return const_cast<ConstBundlePrim*>(this)->path();
}

inline omni::graph::core::NameToken ConstBundlePrim::type() noexcept
{
    using namespace omni::graph::core;

    ConstBundlePrims* bundlePrims = getConstBundlePrims();
    BundlePrimIndex const primIndex = this->primIndex();
    if (primIndex >=  bundlePrims->getPrimCount())
    {
        return carb::flatcache::kUninitializedToken;
    }
    NameToken const* types = bundlePrims->getConstPrimTypes();
    return (types != nullptr) ? types[primIndex] : carb::flatcache::kUninitializedToken;
}

inline omni::graph::core::NameToken ConstBundlePrim::type() const noexcept
{
    return const_cast<ConstBundlePrim*>(this)->type();
}

inline DirtyIDType ConstBundlePrim::dirtyID() noexcept
{
    ConstBundlePrims* bundlePrims = getConstBundlePrims();
    if (primIndex() >= bundlePrims->getPrimCount())
    {
        return bundlePrims->getBundleDirtyID();
    }

    DirtyIDType const* dirtyIDs = m_bundlePrims->getPrimDirtyIDs();
    return (dirtyIDs != nullptr) ? dirtyIDs[m_primIndex] : kInvalidDirtyID;
}

inline DirtyIDType ConstBundlePrim::dirtyID() const noexcept
{
    return const_cast<ConstBundlePrim*>(this)->dirtyID();
}

inline ConstBundlePrims* ConstBundlePrim::getConstBundlePrims() noexcept
{
    return m_bundlePrims;
}

inline ConstBundlePrimAttrIterator ConstBundlePrim::begin() noexcept
{
    return ConstBundlePrimAttrIterator(*this, getAttributes().begin());
}

inline ConstBundlePrimAttrIterator ConstBundlePrim::end() noexcept
{
    return ConstBundlePrimAttrIterator(*this, getAttributes().end());
}

inline ConstBundlePrimAttrIterator ConstBundlePrim::begin() const noexcept
{
    ConstBundlePrim& thisPrim = const_cast<ConstBundlePrim&>(*this);
    return ConstBundlePrimAttrIterator(thisPrim, thisPrim.getAttributes().begin());
}

inline ConstBundlePrimAttrIterator ConstBundlePrim::end() const noexcept
{
    ConstBundlePrim& thisPrim = const_cast<ConstBundlePrim&>(*this);
    return ConstBundlePrimAttrIterator(thisPrim, thisPrim.getAttributes().end());
}

inline omni::graph::core::IConstBundle2* ConstBundlePrim::getConstBundlePtr() noexcept
{
    return m_bundle.get();
}

inline ConstBundlePrim::BundleAttributeMap& ConstBundlePrim::getAttributes() noexcept
{
    return m_attributes;
}

// ====================================================================================================
//
//                                      Const Bundle Primitives
//
// ====================================================================================================

inline ConstBundlePrims::ConstBundlePrims()
{
}

inline ConstBundlePrims::ConstBundlePrims(omni::graph::core::GraphContextObj const& context,
                                          omni::graph::core::ConstBundleHandle const& bundle)
    : ConstBundlePrims()
{
    attach(context, bundle);
}

inline void ConstBundlePrims::detach() noexcept
{
    m_bundleDirtyID = kInvalidDirtyID;
    m_primTypes = nullptr;
    m_primPaths = nullptr;
    m_primDirtyIDs = nullptr;
    m_iDirtyID = nullptr;
    m_primitives.clear();
    m_commonAttributes.reset();
    m_context = omni::graph::core::GraphContextObj{};
    m_bundle.release();
    m_factory.release();
}

inline omni::graph::core::NameToken const* ConstBundlePrims::getConstPrimPaths() noexcept
{
    return m_primPaths;
}

inline ConstBundlePrims::BundlePrimArray& ConstBundlePrims::getPrimitives() noexcept
{
    return m_primitives;
}

inline omni::graph::core::NameToken const* ConstBundlePrims::getConstPrimTypes() noexcept
{
    return m_primTypes;
}

inline omni::graph::core::ConstBundleHandle ConstBundlePrims::getConstHandle() noexcept
{
    return m_bundle->getConstHandle();
}

template <typename FUNC>
ConstBundlePrim* ConstBundlePrims::getConstPrim(BundlePrimIndex primIndex, FUNC createBundlePrim) noexcept
{
    // Return invalid const bundle prim if out of bounds.
    size_t const bundlePrimCount = getPrimCount();
    if (primIndex >= bundlePrimCount)
    {
        return nullptr;
    }

    // Search and return if in cache.
    auto& prims = getPrimitives();
    if (prims.size() != bundlePrimCount)
    {
        prims.resize(bundlePrimCount);
    }

    if (prims[primIndex] != nullptr)
    {
        return prims[primIndex].get();
    }

    // update the cache and return the bundle prim.
    std::unique_ptr<ConstBundlePrim> newBundlePrim{ createBundlePrim() };
    if (!newBundlePrim)
    {
        return nullptr;
    }

    ConstBundlePrim* newBundlePrimPtr = newBundlePrim.get();
    prims[primIndex] = std::move(newBundlePrim);
    return newBundlePrimPtr;
}

inline ConstBundlePrim* ConstBundlePrims::getPrim(BundlePrimIndex primIndex) noexcept
{
    return getConstPrim(primIndex);
}

inline ConstBundlePrim* ConstBundlePrims::getConstPrim(BundlePrimIndex primIndex) noexcept
{
    using namespace omni::graph::core;

    auto createBundlePrim = [this, &bundlePrims = *this, &primIndex]() -> ConstBundlePrim*
    {
        ConstBundleHandle bundleHandle = getConstBundlePtr()->getConstChildBundle(primIndex);
        if (!bundleHandle.isValid())
        {
            return nullptr;
        }

        omni::core::ObjectPtr<IConstBundle2> childBundle = getBundleFactoryPtr()->getConstBundle(context(), bundleHandle);
        if (!childBundle)
        {
            return nullptr;
        }

        return new ConstBundlePrim{ bundlePrims, primIndex, childBundle };
    };

    return getConstPrim(primIndex, createBundlePrim);
}

inline DirtyIDType ConstBundlePrims::getBundleDirtyID() noexcept
{
    return m_bundleDirtyID;
}

inline DirtyIDType const* ConstBundlePrims::getPrimDirtyIDs() noexcept
{
    return m_primDirtyIDs;
}

inline ConstBundlePrim& ConstBundlePrims::getConstCommonAttrs() noexcept
{
    return *m_commonAttributes;
}

inline omni::graph::core::GraphContextObj const& ConstBundlePrims::context() noexcept
{
    using namespace omni::graph::core;

    if (m_bundle)
    {
        m_context = m_bundle->getContext();
    }
    else
    {
        m_context = GraphContextObj{};
    }
    return m_context;
}

inline DirtyIDType ConstBundlePrims::getNextDirtyID() noexcept
{
    return m_iDirtyID->getNextDirtyID();
}

inline void ConstBundlePrims::attach(omni::graph::core::GraphContextObj const& context,
                                     omni::graph::core::ConstBundleHandle const& bundleHandle) noexcept
{
    using namespace omni::graph::core;

    ComputeGraph* computeGraph = carb::getCachedInterface<ComputeGraph>();
    omni::core::ObjectPtr<IBundleFactory> factory = computeGraph->getBundleFactoryInterfacePtr();
    omni::core::ObjectPtr<IConstBundle2> bundle = factory->getConstBundle(context, bundleHandle);
    attach(std::move(factory), std::move(bundle));
}

inline void ConstBundlePrims::attach(omni::core::ObjectPtr<omni::graph::core::IBundleFactory>&& factoryPtr,
                                     omni::core::ObjectPtr<omni::graph::core::IConstBundle2>&& bundlePtr) noexcept
{
    using namespace omni::graph::core;

    // Initialize members
    m_factory = std::move(factoryPtr);
    m_bundle = std::move(bundlePtr);

    // Initialize common attributes to provide access to ConstBundlePrims attributes.
    m_commonAttributes.reset(new ConstBundlePrim(*this, kInvalidBundlePrimIndex, m_bundle));

    // Acquire IDirtyID interface
    m_iDirtyID = carb::getCachedInterface<omni::graph::io::IDirtyID>();

    if (!m_bundle->isValid())
    {
        return;
    }

    // TODO: Following code is necessary for backward compatibility.
    IConstBundle2* bundle = getConstBundlePtr();
    GraphContextObj const& context = this->context();

    // Bundle DirtyID.
    auto& bundleDirtyIDDef = detail::getBundleDirtyIDDefinition();
    ConstAttributeDataHandle bundleDirtyIDHandle = bundle->getConstBundleMetadataByName(bundleDirtyIDDef.token);
    if (bundleDirtyIDHandle.isValid())
    {
        setBundleDirtyID(*getDataR<DirtyIDType>(context, bundleDirtyIDHandle));
    }
    else
    {
        setBundleDirtyID(kInvalidDirtyID);
    }

    // Prim DirtyIDs.
    auto& primDirtyIDsDef = detail::getPrimDirtyIDsDefinition();
    ConstAttributeDataHandle primDirtyIDsHandle = bundle->getConstBundleMetadataByName(primDirtyIDsDef.token);
    if (primDirtyIDsHandle.isValid())
    {
        size_t arrayLength = 0;
        context.iAttributeData->getElementCount(&arrayLength, context, &primDirtyIDsHandle, 1);
        setPrimDirtyIDsData(*getDataR<DirtyIDType*>(context, primDirtyIDsHandle));
    }
    else
    {
        setPrimDirtyIDsData(nullptr);
    }

    // Prim Paths.
    auto& primPathsDef = detail::getPrimPathsDefinition();
    ConstAttributeDataHandle primPathsHandle = bundle->getConstBundleMetadataByName(primPathsDef.token);
    if (primPathsHandle.isValid())
    {
        size_t arrayLength = 0;
        context.iAttributeData->getElementCount(&arrayLength, context, &primPathsHandle, 1);
        setPrimPathsData(*getDataR<NameToken*>(context, primPathsHandle));
    }
    else
    {
        setPrimPathsData(nullptr);
    }

    // Prim Types.
    auto& primTypesDef = detail::getPrimTypesDefinition();
    ConstAttributeDataHandle primTypesHandle = bundle->getConstBundleMetadataByName(primTypesDef.token);
    if (primTypesHandle.isValid())
    {
        size_t arrayLength = 0;
        context.iAttributeData->getElementCount(&arrayLength, context, &primTypesHandle, 1);
        setPrimTypesData(*getDataR<NameToken*>(context, primTypesHandle));
    }
    else
    {
        setPrimTypesData(nullptr);
    }
}

inline void ConstBundlePrims::setBundleDirtyID(DirtyIDType bundleDirtyID) noexcept
{
    m_bundleDirtyID = bundleDirtyID;
}

inline void ConstBundlePrims::setPrimDirtyIDsData(DirtyIDType const* primDirtyIDs) noexcept
{
    m_primDirtyIDs = primDirtyIDs;
}

inline void ConstBundlePrims::setPrimPathsData(omni::graph::core::NameToken const* primPaths) noexcept
{
    m_primPaths = primPaths;
}

inline void ConstBundlePrims::setPrimTypesData(omni::graph::core::NameToken const* primTypes) noexcept
{
    m_primTypes = primTypes;
}

inline omni::graph::core::IBundleFactory* ConstBundlePrims::getBundleFactoryPtr() noexcept
{
    return m_factory.get();
}

inline omni::graph::core::IConstBundle2* ConstBundlePrims::getConstBundlePtr() noexcept
{
    return m_bundle.get();
}

inline size_t ConstBundlePrims::getPrimCount() noexcept
{
    if (IConstBundle2* bundle = getConstBundlePtr())
    {
        return bundle->getChildBundleCount();
    }
    return 0;
}

inline ConstBundlePrimIterator ConstBundlePrims::begin() noexcept
{
    return ConstBundlePrimIterator(*this);
}

inline ConstBundlePrimIterator ConstBundlePrims::end() noexcept
{
    return ConstBundlePrimIterator(*this, getPrimCount());
}

/***********************************************************************************************
 *
 * TODO: Following methods might be deprecated in the future, but are kept for backward compatibility.
 *       In the next iteration when real interface starts to emerge, we can retire those methods.
 *
 ***********************************************************************************************/

inline ConstBundlePrim& ConstBundlePrims::getCommonAttrs() noexcept
{
    return getConstCommonAttrs();
}

inline omni::graph::core::ConstBundleHandle ConstBundlePrims::handle() noexcept
{
    return m_bundle->getConstHandle();
}

inline omni::graph::core::NameToken const* ConstBundlePrims::getPrimPaths() noexcept
{
    return getConstPrimPaths();
}

inline void ConstBundlePrims::separateAttrs() noexcept
{
    // There is nothing to separate. This function is deprecated.
}

inline void ConstBundlePrims::ensurePrimAttrsCached(BundlePrimIndex primIndex) noexcept
{
    // Responsibility of caching attributes was moved to Bundle Prim.
}

// ====================================================================================================
//
//                                      Const Bundle Primitive Iterator
//
// ====================================================================================================

inline ConstBundlePrimIterator::ConstBundlePrimIterator(ConstBundlePrims& bundlePrims, BundlePrimIndex primIndex) noexcept
    : m_bundlePrims(&bundlePrims), m_primIndex(primIndex)
{
}

inline bool ConstBundlePrimIterator::operator==(ConstBundlePrimIterator const& that) const noexcept
{
    return m_bundlePrims == that.m_bundlePrims && m_primIndex == that.m_primIndex;
}

inline bool ConstBundlePrimIterator::operator!=(ConstBundlePrimIterator const& that) const noexcept
{
    return !(*this == that);
}

inline ConstBundlePrim& ConstBundlePrimIterator::operator*() noexcept
{
    return *(m_bundlePrims->getConstPrim(m_primIndex));
}

inline ConstBundlePrim* ConstBundlePrimIterator::operator->() noexcept
{
    return m_bundlePrims->getConstPrim(m_primIndex);
}

inline ConstBundlePrimIterator& ConstBundlePrimIterator::operator++() noexcept
{
    ++m_primIndex;
    return *this;
}

// ====================================================================================================
//
//                                      Const Bundle Primitive Attribute Iterator
//
// ====================================================================================================

inline ConstBundlePrimAttrIterator::ConstBundlePrimAttrIterator(ConstBundlePrim& bundlePrim, ConstBundlePrim::AttrMapIteratorType attrIter) noexcept
    : m_bundlePrim(&bundlePrim), m_attrIter(attrIter)
{
}

inline bool ConstBundlePrimAttrIterator::operator==(ConstBundlePrimAttrIterator const& that) const noexcept
{
    return m_bundlePrim == that.m_bundlePrim && m_attrIter == that.m_attrIter;
}

inline bool ConstBundlePrimAttrIterator::operator!=(ConstBundlePrimAttrIterator const& that) const noexcept
{
    return !(*this == that);
}

inline BundleAttrib const& ConstBundlePrimAttrIterator::operator*() const noexcept
{
    CARB_ASSERT(m_attrIter->second);
    return *(m_attrIter->second);
}

inline BundleAttrib const*ConstBundlePrimAttrIterator:: operator->() const noexcept
{
    CARB_ASSERT(m_attrIter->second);
    return m_attrIter->second.get();
}

inline ConstBundlePrimAttrIterator& ConstBundlePrimAttrIterator::operator++() noexcept
{
    ++m_attrIter;
    return *this;
}

} // namespace io
} // namespace graph
} // namespace omni
