/******************************************************************************
 * File: Handles.hpp
 *
 * Copyright (c) 2023 Qualcomm Technologies, Inc.
 * All Rights Reserved
 * Confidential and Proprietary - Qualcomm Technologies, Inc.
 *
 ******************************************************************************/

#ifndef QCOM_HANDLES_HPP
#define QCOM_HANDLES_HPP

#include <cassert>
#include <list>
#include <mutex>
#include <unordered_map>

#include "openxr/openxr.h"

namespace wikitude::openxr {

    // Thread-safe OpenXR handle allocator.
    // This is a singleton object as OpenXr handles must be unique process-wide.
    class HandleAllocator {
    public:
        static HandleAllocator& get() {
            return g_instance;
        }

        template<typename T>
        T createXrHandle() {
            std::unique_lock _{ _mutex };

            if (_availableHandles.size() == 0) {
                growAvailableHandlesContainer();    // NOTE(DG): create new handles if all the existing ones have been exhausted
            }

            auto it = _availableHandles.begin();
            T handle = reinterpret_cast<T>(*it);
            _availableHandles.erase(it);

            return handle;
        }

        template<typename T>
        void destroyXrHandle(T handle_) {
            std::unique_lock _{ _mutex };

            _availableHandles.push_front(reinterpret_cast<std::uint64_t>(handle_));    // NOTE(DG): recycle handle value by appending it at the front of the list
        }

    private:
        HandleAllocator();
        HandleAllocator(const HandleAllocator&) = delete;
        HandleAllocator& operator=(const HandleAllocator&) = delete;

        void growAvailableHandlesContainer();

    private:
        static HandleAllocator g_instance;

        std::mutex _mutex;
        std::list<std::uint64_t> _availableHandles;
        std::uint64_t _nextAvailableHandle;
        std::size_t _currentAvailableHandleGrowth;
    };

    // Thread-safe container that maps OpenXR handles to their equivalent C++ objects.

    template<typename THandle, typename TObject>
    class HandleMap {
    public:
        using HandleType = void*;    // This type has the same size as handles defined with XR_DEFINE_HANDLE()

        static_assert(sizeof(THandle) == sizeof(HandleType), "Template parameter THandle must be a handle type defined with XR_DEFINE_HANDLE()");

        HandleMap() = default;
        HandleMap(const HandleMap&) = delete;
        HandleMap& operator=(const HandleMap&) = delete;

        void insert(THandle handle_, TObject& object_) {
            std::lock_guard _{ _mutex };
            _handleToObject[handle_] = &object_;
        }

        bool erase(THandle handle_) {
            std::lock_guard _{ _mutex };
            size_t erased = _handleToObject.erase(handle_);
            return erased > 0;
        }

        TObject* find(THandle handle_) {
            std::lock_guard _{ _mutex };
            auto itFind = _handleToObject.find(handle_);
            if (itFind != _handleToObject.end()) {
                return static_cast<TObject*>(itFind->second);
            }
            return nullptr;
        }

    private:
        std::mutex _mutex;
        std::unordered_map<HandleType, TObject*> _handleToObject;
    };

    // Wrapper class for a handle type declared with XR_DEFINE_HANDLE().
    // Manages the insertion/removal of the handle and its object into a handle map.
    // Optionnally manages the allocation/deallocation of the handle itself.

    template<typename THandle, typename TObject>
    class HandleWrapper {
    public:
        using THandleMap = HandleMap<THandle, TObject>;

        struct Allocator {
            std::function<THandle(void)> allocate;
            std::function<void(THandle)> deallocate;
        };

        // Construct an HandleWrapper object charged with allocating/deallocating the handle
        // with the default handle allocator.
        HandleWrapper(THandleMap& handleMap_, TObject& object_)
            : _handle(HandleAllocator::get().createXrHandle<THandle>())
            , _handleMap(handleMap_)
            , _ownHandle(true) {
            if (_handle != XR_NULL_HANDLE) {
                _handleMap.insert(_handle, object_);
            }
        }

        // Construct an HandleWrapper object charged with allocating/deallocating the handle
        // using a custom allocator passed as parameter.
        HandleWrapper(THandleMap& handleMap_, TObject& object_, Allocator customAllocator_)
            : _handleMap(handleMap_)
            , _ownHandle(true)
            , _customAllocator(customAllocator_) {
            _handle = _customAllocator->allocate();

            if (_handle != XR_NULL_HANDLE) {
                _handleMap.insert(_handle, object_);
            }
        }

        // Construct an HandleWrapper object from a previously-allocated, externally-managed handle.
        HandleWrapper(THandleMap& handleMap_, TObject& object_, THandle handle_)
            : _handle(handle_)
            , _handleMap(handleMap_)
            , _ownHandle(false) {
            if (_handle != XR_NULL_HANDLE) {
                _handleMap.insert(_handle, object_);
            }
        }

        ~HandleWrapper() {
            static_assert(sizeof(THandle) == sizeof(typename THandleMap::HandleType), "Template parameter THandle must be a handle type defined with XR_DEFINE_HANDLE()");

            if (_handle != XR_NULL_HANDLE) {
                _handleMap.erase(_handle);

                if (_ownHandle) {
                    if (_customAllocator) {
                        _customAllocator->deallocate(_handle);
                    } else {
                        HandleAllocator::get().destroyXrHandle<THandle>(_handle);
                    }
                }
            }
        }

        HandleWrapper(const HandleWrapper&) = delete;
        HandleWrapper& operator=(const HandleWrapper&) = delete;

        operator THandle() const {
            return get();
        }

        THandle get() const {
            return _handle;
        }

    private:
        THandle _handle;
        THandleMap& _handleMap;
        bool _ownHandle;
        std::optional<Allocator> _customAllocator;
    };

    // RAII helper for locking an instance-scoped mutex from any of that instance's child handles.
    //
    // This class ensures that the instance mutex is locked for the entire duration
    // of a call into a C++ OpenXR object. This helps guarantee that access to an instance C++ object
    // and to all of its subobjects (sessions, trackers, targets...) is single threaded.
    //
    // This class performs the following operations:
    // - look for a specific handle in a HandleMap
    // - if the OpenXR object is found in the map, retrieve its parent instance object
    // - lock the instance mutex
    // - hold the lock until the class destructor.

    template<typename THandle, typename TObject>
    class InstanceObjectLock {
    public:
        using InstanceObjectType = typename TObject::InstanceObjectType;

        static_assert(sizeof(THandle) == sizeof(typename HandleMap<THandle, TObject>::HandleType), "Template parameter THandle must be a handle type defined with XR_DEFINE_HANDLE()");

        InstanceObjectLock(HandleMap<THandle, TObject>& handleMap_, THandle handle_)
            : _instance(nullptr) {
            _object = handleMap_.find(handle_);
            if (_object) {
                _instance = &_object->getInstance();
                _instance->lock();
            }
        }

        InstanceObjectLock(const InstanceObjectLock&) = delete;
        InstanceObjectLock& operator=(const InstanceObjectLock&) = delete;

        ~InstanceObjectLock() {
            if (_instance) {
                _instance->unlock();
            }
        }

        operator bool() {
            return _object != nullptr;
        }

        TObject& getObject() {
            assert(bool(*this));
            return *_object;
        }

        InstanceObjectType& getInstance() {
            assert(bool(*this));
            return *_instance;
        }

    private:
        TObject* _object;
        InstanceObjectType* _instance;
    };
}

#endif /* QCOM_HANDLES_HPP */
