// Copyright (C) 2018-2020 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

/**
 * @brief A header file for generic LockedMemory<> and different variations of locks
 *
 * @file ie_locked_memory.hpp
 */
#pragma once

#include <iostream>
#include <utility>

#include "ie_allocator.hpp"

namespace InferenceEngine {
namespace details {
/**
 * @brief This class is a LockedMemory concept for hardware memory
 */
template <class T>
class LockedMemoryBase {
    IAllocator* _allocator = nullptr;
    void* _handle = nullptr;
    mutable T* _locked = nullptr;
    LockOp _lockFlag = LOCK_FOR_WRITE;

protected:
    /**
     * @brief An offset size
     *
     * The default value is 0.
     */
    size_t _offset = 0;

public:
    /**
     * @brief A constructor
     *
     * @param ptr Pointer to an IAllocator object
     * @param handle Handle provided by allocator->Alloc()
     * @param lockFlag Read/Write type of mapping
     * @param offsetInBytes Offset in originally locked region
     */
    LockedMemoryBase(IAllocator* ptr, void* handle, LockOp lockFlag, size_t offsetInBytes)
        : _allocator(ptr), _handle(handle), _lockFlag(lockFlag), _offset(offsetInBytes) {}

    /**
     * @brief A copy constructor
     *
     * @param that An rvalue reference for the other LockedMemoryBase instance
     */
    LockedMemoryBase(LockedMemoryBase&& that)
        : _allocator(that._allocator), _handle(that._handle), _lockFlag(that._lockFlag), _offset(that._offset) {
        that._locked = nullptr;
    }

    /**
     * @brief A virtual destructor
     */
    virtual ~LockedMemoryBase() {
        if (_locked != nullptr) {
            _allocator->unlock(_handle);
        }
    }

protected:
    /**
     * @brief Compares referenced values
     *
     * @param pointer Pointer to the object to compare with
     * @return True if all handlers are nullptr or referenced values are equal, false otherwise
     */
    bool isEqualTo(const T* pointer) const {
        if (pointer == nullptr && (_allocator == nullptr || _handle == nullptr)) {
            return true;
        }
        return dereference() == pointer;
    }

    /**
     * @brief Gets the locked object.
     *
     * Locks the handler and casts memory to the object of the given template type.
     *
     * @return The pointer to the locked object, nullptr otherwise
     */
    virtual T* dereference() const {
        if (_locked != nullptr) return _locked;

        if (_allocator == nullptr) {
            return nullptr;
        }

        if (_handle == nullptr) {
            return nullptr;
        }

        uint8_t* pBytes = reinterpret_cast<uint8_t*>(_allocator->lock(_handle, _lockFlag));

        return _locked = reinterpret_cast<T*>(pBytes + _offset);
    }
};
}  // namespace details

/**
 * @brief This class represents locked memory for read/write memory
 */
template <class T>
class LockedMemory : public details::LockedMemoryBase<T> {
    using base = details::LockedMemoryBase<T>;

public:
    /**
     * @brief A constructor
     *
     * @param ptr Pointer to IAllocator object
     * @param handle Handle provided by allocator
     * @param offsetInBytes Offset in originally locked region
     */
    LockedMemory(IAllocator* ptr, void* handle, size_t offsetInBytes = 0)
        : base(ptr, handle, LOCK_FOR_WRITE, offsetInBytes) {}

    /**
     * @brief A default copy constructor, accepting rvalue
     */
    LockedMemory(LockedMemory<T>&&) = default;

    /**
     * @brief A default copy constructor that accepts rvalue
     *
     * Also sets the offset value for the new memory object
     *
     * @param that Rvalue reference for the other LockedMemoryBase instance
     * @param offset Offset value
     */
    LockedMemory(LockedMemory<T>&& that, size_t offset): base(std::move(that)) {
        base::_offset = offset;
    }

    /**
     * @brief A disabled copy constructor for lvalue
     */
    LockedMemory(const LockedMemory<T>&) = delete;

    /**
     * @brief Gets a pointer to the stored object
     *
     * Dereferences from the base class.
     *
     * @return The pointer to the object of the given template type
     */
    operator T*() {
        return base::dereference();
    }

    /**
     * @brief Gets the const pointer to the stored object
     *
     * Dereferences from the base class.
     * @return The const pointer object of the given template type.
     */
    operator const T*() const {
        return base::dereference();
    }

    /**
     * @brief Compares stored object with the given one
     * @param pointer An pointer to compare with.
     * @return true if objects are equal, false otherwise
     */
    bool operator==(const T* pointer) const {
        // special case with nullptr
        return base::isEqualTo(pointer);
    }

    /**
     * @brief Compares the object with the one stored in the memory.
     *
     * @return true if objects are equal, false otherwise
     */
    friend bool operator==(const T* pointer, const LockedMemory<T>& lm) {
        return lm.operator==(pointer);
    }

    /**
     * @brief Casts stored object to any provided type.
     *
     * Uses reinterpret_cast.
     *
     * @tparam S Type to be casted to
     * @return Casted to the given type object
     */
    template <class S, typename = std::enable_if<std::is_pointer<S>::value>>
    S as() {
        return reinterpret_cast<S>(base::dereference());
    }

    /**
     * @brief Casts stored object to any provided type.
     *
     * Uses reinterpret_cast.
     *
     * @tparam S Type to be casted to
     * @return Casted to the given type const object
     */
    template <class S, typename = std::enable_if<std::is_pointer<S>::value>>
    const S as() const {
        return reinterpret_cast<S>(base::dereference());
    }
};

/**
 * @brief This class is for <void*> data and allows casting to any pointers
 */
template <>
class LockedMemory<void> : public details::LockedMemoryBase<void> {
    using base = details::LockedMemoryBase<void>;

public:
    /**
     * @brief A constructor
     *
     * @param ptr Pointer to IAllocator object
     * @param handle Handle provided by allocator
     * @param offsetInBytes Offset in originally locked region
     */
    LockedMemory(IAllocator* ptr, void* handle, size_t offsetInBytes)
        : base(ptr, handle, LOCK_FOR_WRITE, offsetInBytes) {}

    /**
     * @brief A default copy constructor that accepts rvalue
     */
    LockedMemory(LockedMemory<void>&&) = default;

    /**
     * @brief A default copy constructor that accepts rvalue
     *
     * Also sets the offset value for the new memory object
     *
     * @param that Rvalue reference for the other LockedMemoryBase instance
     * @param offset Offset value
     */
    LockedMemory(LockedMemory<void>&& that, size_t offset): base(std::move(that)) {
        base::_offset = offset;
    }

    /**
     * @brief A disabled copy constructor for lvalue
     */
    LockedMemory(const LockedMemory<void>&) = delete;

    /**
     * @brief Gets the pointer to the stored object of the given template type
     *
     * Dereferences from the base class.
     *
     * @tparam S Type to be casted to
     * @return The pointer to the object of the given template type
     */
    template <class S>
    operator S*() {
        return reinterpret_cast<S*>(base::dereference());
    }

    /**
     * @brief Compares stored object with the given one
     *
     * @return true if objects are equal, false otherwise
     */
    bool operator==(const void* pointer) const {
        // special case with nullptr
        return base::isEqualTo(pointer);
    }

    /**
     * @brief Compares the object with the one stored in the memory
     *
     * @return true if objects are equal, false otherwise
     */
    friend bool operator==(const void* pointer, const LockedMemory<void>& lm) {
        return lm.operator==(pointer);
    }

    /**
     * @brief Casts stored object to any given type
     *
     * Uses reinterpret_cast.
     *
     * @tparam S Type to be casted to
     * @return Casted to the given type object
     */
    template <class S, typename = std::enable_if<std::is_pointer<S>::value>>
    S as() {
        return reinterpret_cast<S>(dereference());
    }

    /**
     * @brief Casts stored object to any given type
     *
     * Uses reinterpret_cast.
     *
     * @tparam S Type to be casted to
     * @return Casted to the given type const object
     */
    template <class S, typename = std::enable_if<std::is_pointer<S>::value>>
    const S as() const {
        return reinterpret_cast<S>(dereference());
    }
};

/**
 * @brief This class is for read-only segments
 */
template <class T>
class LockedMemory<const T> : public details::LockedMemoryBase<T> {
    using base = details::LockedMemoryBase<T>;

public:
    /**
     * @brief A constructor
     *
     * @param ptr Pointer to IAllocator object
     * @param handle Handle provided by allocator
     * @param offset Offset in bytes in originally locked region
     */
    LockedMemory(IAllocator* ptr, void* handle, size_t offset): base(ptr, handle, LOCK_FOR_READ, offset) {}

    /**
     * @brief A default copy constructor that accepts rvalue
     */
    LockedMemory(LockedMemory<const T>&&) = default;

    /**
     * @brief A default copy constructor that accepts rvalue.
     *
     * Also sets the offset value for the new memory object
     *
     * @param that Rvalue reference for the other LockedMemoryBase instance
     * @param offset Offset value
     */
    LockedMemory(LockedMemory<const T>&& that, size_t offset): base(std::move(that)) {
        base::_offset = offset;
    }

    /**
     * @brief A disabled copy constructor for lvalue
     */
    LockedMemory(const LockedMemory<const T>&) = delete;

    /**
     * @brief Gets the const pointer to the stored object
     *
     * Dereferences from the base class.
     *
     * @return The pointer to the object.
     */
    operator const T*() const {
        return base::dereference();
    }

    /**
     * @brief Compares stored object with the given one
     *
     * @return true if objects are equal, false otherwise
     */
    bool operator==(const T* pointer) const {
        // special case with nullptr
        return base::isEqualTo(pointer);
    }

    /**
     * @brief Compares the object with the one stored in the memory
     *
     * @return true if objects are equal, false otherwise
     */
    friend bool operator==(const T* pointer, const LockedMemory<const T>& lm) {
        return lm.operator==(pointer);
    }

    /**
     * @brief Casts stored object to any given type.
     *
     * Uses reinterpret_cast.
     *
     * @tparam S Type to be casted to
     * @return Casted to the given type object
     */
    template <class S, typename = std::enable_if<std::is_pointer<S>::value && std::is_const<S>::value>>
    S as() const {
        return reinterpret_cast<S>(base::dereference());
    }
};
}  // namespace InferenceEngine
