#pragma once

#include <functional>
#include <memory>
#include <unordered_map>

namespace hdk::grid {
/**
 * @brief Reference-counted wrapper for shared pointers to structures.
 * @remarks
 * Has no default constructor. Must be constructed with a value or explicitly set initialize to nullptr.
 * Evaluates to false if shared pointer is null, true otherwise.
 * Will automatically cast to the `StructType*`, allowing it to be used in C API calls that expect raw pointers, while still benefiting from the reference counting and caching mechanisms provided by the wrapper.
 */
template <typename StructType>
class SharedPtrWrapper {
public:
    /** Type alias for a shared pointer to the managed structure. **/
    using SharedPtr = std::shared_ptr<StructType>;
    /** Type alias for a weak pointer to the managed structure. **/
    using WeakPtr = std::weak_ptr<StructType>;
    /** Type alias for the cache map that maps raw pointers to weak pointers. **/
    using SharedCacheMap = std::unordered_map<StructType*, WeakPtr>;
    /** Function type for deleting the managed structure. **/
    using DeleteFunc = std::function<void(StructType*)>;
    static void null_deleter(StructType*) { }
    /** Default constructor deleted. You must declare a value or declare it to be null on construction. **/
    SharedPtrWrapper() = delete;
    /** Null constructor (same as default constructor). **/
    SharedPtrWrapper(std::nullptr_t)
        : shared_struct(nullptr)
    {
    }
    /** Constructor that takes a shared pointer and a delete function. **/
    SharedPtrWrapper(StructType* struct_ptr, DeleteFunc deleteFunc)
        : shared_struct(struct_ptr, deleteFunc)
    {
        // Cache the shared pointer to ensure consistent reference counting for the same raw pointer.
        get_shared_weak_ptr_cache()[struct_ptr] = shared_struct;
    }
    /** Move constructor. **/
    SharedPtrWrapper(SharedPtrWrapper&& other) noexcept
        : shared_struct(std::move(other.shared_struct))
    {
    }
    /** Existing Shared Pointer */
    SharedPtrWrapper(SharedPtr shared_ptr)

    {
        shared_struct = shared_ptr;
        if ((shared_ptr.get() != nullptr) && !is_cached(shared_ptr.get())) {
            get_shared_weak_ptr_cache()[shared_ptr.get()] = shared_struct;
        }
    }
    /** Move assignment operator. **/
    SharedPtrWrapper& operator=(SharedPtrWrapper&& other) noexcept
    {
        if (this != &other) {
            shared_struct = std::move(other.shared_struct);
        }
        return *this;
    }
    /** Copy constructor. **/
    SharedPtrWrapper(const SharedPtrWrapper& other)
        : shared_struct(other.shared_struct)
    {
    }
    /** Copy assignment operator. **/
    SharedPtrWrapper& operator=(const SharedPtrWrapper& other)
    {
        if (this != &other) {
            shared_struct = other.shared_struct;
        }
        return *this;
    }
    /** Conversion operator to allow automatic casting to StructType*. **/
    operator StructType*() const
    {
        return shared_struct.get();
    }
    /** Set this instance to nullptr. **/
    SharedPtrWrapper& operator=(std::nullptr_t)
    {
        shared_struct.reset();
        return *this;
    }
    /** Conversion operator to allow automatic casting to bool. **/
    operator bool() const
    {
        return shared_struct != nullptr;
    }

private: // internal use only.
    /** The value we are wrapping. **/
    SharedPtr shared_struct{nullptr};

protected: // For use by derived classes.
    /** So that we provide consistent reference counting on the same pointer we need to have a cache of shared pointers. This is because if we create multiple shared pointers to the same raw pointer, they will have separate reference counts and may lead to double deletion. **/
    static SharedCacheMap& get_shared_weak_ptr_cache()
    {
        static SharedCacheMap cache_map;
        return cache_map;
    }
    /** Helper function to check if a raw pointer is already cached and has a valid shared pointer. **/
    static bool is_cached(StructType* ptr)
    {
        auto& cache = get_shared_weak_ptr_cache();
        return cache.find(ptr) != cache.end() && !cache[ptr].expired();
    }
    /** @brief helper: if cached returns existing, else returns non-owning shared pointer without caching it.
     * @remarks
     * The idea is that if we are already managing this pointer
     * we should return the existing shared pointer to ensure consistent reference counting
     * but if we are not managing this pointer, we can create a new shared pointer with a null deleter
     * so that we can still return a shared pointer without taking ownership of the pointer.
     * @return Always a valid shared pointer to `ptr`, but it may be a non-owning shared pointer if `ptr` is not already cached.
     ***/
    static SharedPtr get_or_view(StructType* ptr)
    {
        /** if no pointer */
        if (!ptr) {
            /** then return a non-owning null */
            return SharedPtr(nullptr, null_deleter);
        }

        if (is_cached(ptr)) {
            return get_shared_weak_ptr_cache()[ptr].lock();
        } else {
            SharedPtr new_shared(ptr, null_deleter);
            // cache[ptr] = new_shared;
            return new_shared;
        }
    }
    /** @brief helper: if cached returns existing, else constructs shared pointer with deleter and caches.
     * @remarks
     * The idea is that if we are already managing this pointer
     * we should return the existing shared pointer to ensure consistent reference counting
     * but if we are not managing this pointer, we can create a new shared pointer taking ownership.
     * When this pointer is given in the future, as long as it is still valid, will return the same shared pointer.
     * @param ptr the pointer to either get or create for
     * @param deleteFunc the delete function used if creating
     * @return SharedPtr either precached or created and cached
     */
    static SharedPtr get_or_cache(StructType* ptr, DeleteFunc deleteFunc)
    {
        /** If no ptr, no need to do anything else  */
        if (!ptr) {
            return SharedPtr(nullptr, deleteFunc);
        }
        auto& cache = get_shared_weak_ptr_cache();
        if (is_cached(ptr)) {
            return cache[ptr].lock();
        } else {
            SharedPtr new_shared(ptr, deleteFunc);
            cache[ptr] = new_shared;
            return new_shared;
        }
    }
};
}