////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2026 OVITO GmbH, Germany
//
//  This file is part of OVITO (Open Visualization Tool).
//
//  OVITO is free software; you can redistribute it and/or modify it either under the
//  terms of the GNU General Public License version 3 as published by the Free Software
//  Foundation (the "GPL") or, at your option, under the terms of the MIT License.
//  If you do not alter this notice, a recipient may use your version of this
//  file under either the GPL or the MIT License.
//
//  You should have received a copy of the GPL along with this program in a
//  file LICENSE.GPL.txt.  You should have received a copy of the MIT License along
//  with this program in a file LICENSE.MIT.txt
//
//  This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
//  either express or implied. See the GPL or the MIT License for the specific language
//  governing rights and limitations.
//
////////////////////////////////////////////////////////////////////////////////////////

#pragma once

#include <ovito/core/Core.h>

#include <algorithm>

namespace Ovito {

namespace detail {

// Helper alias template
template <typename, typename = std::void_t<>>
struct has_resource_validation_check : std::false_type {};

// Specialization that checks for the existence of `bool isRendererResourceValid()`
template <typename T>
struct has_resource_validation_check<T, std::void_t<decltype(std::declval<T>().isRendererResourceValid())>> : std::true_type {};

} // namespace detail

/**
 * \brief A tagged (=strongly-typed) tuple, which can be used as key for the RendererResourceCache class.
 */
template<typename TagType, typename... TupleFields>
struct RendererResourceKey : public std::tuple<TupleFields...>
{
    /// Inherit constructors of tuple type.
    using std::tuple<TupleFields...>::tuple;
};

/**
 * \brief A cache data structure that accepts keys with arbitrary type and which handles resource lifetime.
 *
 * All methods of the class are thread-safe.
 */
class RendererResourceCache : public std::enable_shared_from_this<RendererResourceCache>
{
    Q_DISABLE_COPY_MOVE(RendererResourceCache)

public:

    /// Data type used by the resource cache to refer to a frame being in flight.
    using ResourceFrameHandle = int;

    /// A frame being in flight.
    class ResourceFrame
    {
    public:
        /// Default constructor.
        ResourceFrame() = default;

        /// Destructor.
        ~ResourceFrame() {
            if(_cache)
                _cache->releaseResourceFrame(_frameNumber);
        }

        // A frame cannot be copied.
        ResourceFrame(const ResourceFrame& other) = delete;
        ResourceFrame& operator=(const ResourceFrame& other) = delete;

        // A frame can be moved.
        ResourceFrame(ResourceFrame&& rhs) noexcept : _cache(std::exchange(rhs._cache, {})), _frameNumber(std::exchange(rhs._frameNumber, 0)) { OVITO_ASSERT(!rhs._cache); }
        ResourceFrame& operator=(ResourceFrame&& rhs) noexcept {
            ResourceFrame(std::move(rhs)).swap(*this);
            return *this;
        }
        inline void swap(ResourceFrame& rhs) noexcept {
            _cache.swap(rhs._cache);
            std::swap(_frameNumber, rhs._frameNumber);
        }

        /// Returns a reference to the cached value for the given key.
        /// Creates a new cache entry with a default-initialized value if the key doesn't exist.
        template<typename Value, typename Key, typename Initializer>
        const Value& lookup(Key&& key, Initializer&& initializer) const {
            OVITO_ASSERT(_cache && _frameNumber != 0);
            return _cache->lookup<Value>(std::forward<Key>(key), _frameNumber, std::forward<Initializer>(initializer));
        }

        /// Returns true if the frame is valid.
        inline operator bool() const noexcept { return (bool)_cache; }

    private:
        /// Constructor.
        ResourceFrame(std::shared_ptr<RendererResourceCache> cache, ResourceFrameHandle frameNumber) noexcept : _cache(std::move(cache)), _frameNumber(frameNumber) {}

        std::shared_ptr<RendererResourceCache> _cache;
        ResourceFrameHandle _frameNumber = 0;

        friend class RendererResourceCache;
    };

public:

    /// Constructor.
    RendererResourceCache() = default;

#ifdef OVITO_DEBUG
    /// Destructor.
    ~RendererResourceCache() {
        // The cache should be completely empty at the time it is destroyed.
        OVITO_ASSERT(_activeResourceFrames.empty());
        OVITO_ASSERT(_entries.empty());
    }
#endif

    /// Returns a reference to the value for the given key.
    /// Creates a new cache entry with a default-initialized value if the key doesn't exist.
    template<typename Value, typename Key, typename Initializer>
    const Value& lookup(Key&& key, ResourceFrameHandle resourceFrame, Initializer&& initializer) {
        std::lock_guard lock(_mutex);
        OVITO_ASSERT(std::find(_activeResourceFrames.begin(), _activeResourceFrames.end(), resourceFrame) != _activeResourceFrames.end());

        // Check if the key exists in the cache.
        for(CacheEntry& entry : _entries) {
            if(entry.key.type() == typeid(Key) && entry.value.type() == typeid(Value) && key == any_cast<const Key&>(entry.key)) {

                // Check at compile time if the Value class has a method 'isRendererResourceValid()'.
                // If so, call it to check if the resource is still valid.
                // This extra check is used, for example, for detecting expired OpenGL textures (see OpenGLTexture class).
                if constexpr(detail::has_resource_validation_check<Value>::value) {
                    if(!any_cast<const Value&>(entry.value).isRendererResourceValid()) {
                        // If the resource is no longer valid, skip this cache entry and create a new one below.
                        continue;
                    }
                }

                // Keep track of the frame in which the requested resource was accessed most recently.
                if(std::ranges::find(entry.frames, resourceFrame) == entry.frames.end()) entry.frames.push_back(resourceFrame);
                // Return reference to the value.
                return any_cast<Value&>(entry.value);
            }
        }

        // Create a new key-value pair with a default-constructed value.
        _entries.emplace_back(std::forward<Key>(key), resourceFrame);
        boost::anys::unique_any& value = _entries.back().value;
        Value& v = value.emplace<Value>();
        // Let the initializer function initialize the value.
        // Note: The initializer may perform more cache lookups, that's why we are using a recursive mutex here.
        if constexpr(std::is_invocable_v<Initializer, Value&>)
            initializer(v);
        else
            std::apply(std::forward<Initializer>(initializer), v);
        OVITO_ASSERT(value.type() == typeid(Value));
        return v;
    }

    /// Opens a new frame with the resource manager.
    ResourceFrame acquireResourceFrame() {
        std::lock_guard lock(_mutex);

        // On the first frame, the cache should be empty.
        OVITO_ASSERT(!_activeResourceFrames.empty() || _entries.empty());

        // Wrap around counter.
        if(_nextResourceFrame == std::numeric_limits<ResourceFrameHandle>::max())
            _nextResourceFrame = 0;

        // Add it to the list of active frames.
        _nextResourceFrame++;
#ifdef OVITO_DEBUG
        _activeResourceFrames.push_back(_nextResourceFrame);
#endif
        return ResourceFrame(shared_from_this(), _nextResourceFrame);
    }

private:

    /// Informs the resource manager that the resources associated with the given frame can be released.
    void releaseResourceFrame(ResourceFrameHandle frame) {
        OVITO_ASSERT(frame > 0);
        std::lock_guard lock(_mutex);

#ifdef OVITO_DEBUG
        // Remove frame from the list of active frames.
        // There is no need to maintain the original list order.
        // We can move the last item into the erased list position.
        auto iter = std::find(_activeResourceFrames.begin(), _activeResourceFrames.end(), frame);
        OVITO_ASSERT(iter != _activeResourceFrames.end());
        *iter = _activeResourceFrames.back();
        _activeResourceFrames.pop_back();
#endif

        // Release the resources associated with the frame unless they are shared with another frame that is still in flight.
        auto end = _entries.end();
        for(auto entry = _entries.begin(); entry != end; ) {
            auto frameIter = std::find(entry->frames.begin(), entry->frames.end(), frame);
            if(frameIter != entry->frames.end()) {
                if(entry->frames.size() != 1) {
                    *frameIter = entry->frames.back();
                    entry->frames.pop_back();
                }
                else {
                    --end;
                    *entry = std::move(*end);
                    continue;
                }
            }
            ++entry;
        }
        _entries.erase(end, _entries.end());

        OVITO_ASSERT(!_activeResourceFrames.empty() || _entries.empty());
    }

    struct CacheEntry {
        template<typename Key> CacheEntry(Key&& _key, ResourceFrameHandle _frame) noexcept : key(std::forward<Key>(_key)) { frames.push_back(_frame); }
        boost::anys::unique_any key;
        boost::anys::unique_any value;
        QVarLengthArray<ResourceFrameHandle, 6> frames;

        // A cache entry cannot be copied.
        CacheEntry(const CacheEntry& other) = delete;
        CacheEntry& operator=(const CacheEntry& other) = delete;

        // A cache entry can be moved.
        CacheEntry(CacheEntry&& other) noexcept : key(std::move(other.key)), value(std::move(other.value)), frames(std::move(other.frames)) {}
        CacheEntry& operator=(CacheEntry&& other) noexcept {
            key = std::move(other.key);
            value = std::move(other.value);
            frames = std::move(other.frames);
            return *this;
        }
    };

    /// Stores all key-value pairs of the cache.
    /// Note we are using std::deque instead of std::vector here, because we require stability of object addresses.
    /// lookup() returns references to elements stored in the cache, which must remain valid even when new objects are added.
    std::deque<CacheEntry> _entries;

    /// Mutex to protect the cache from concurrent accesses.
    std::recursive_mutex _mutex;

    /// Counter that keeps track of how many resource frames have been acquired in total.
    ResourceFrameHandle _nextResourceFrame = 0;

#ifdef OVITO_DEBUG
    /// List of frames that are currently being rendered (by the CPU and/or the GPU).
    std::vector<ResourceFrameHandle> _activeResourceFrames;
#endif
};

}   // End of namespace
