/*
 * Copyright 2022 Haiku Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 */

#ifndef _B_EXCLUSIVE_BORROW_H
#define _B_EXCLUSIVE_BORROW_H

#include <atomic>
#include <memory>

#include <ErrorsExt.h>

namespace BPrivate {

namespace Network {


class BBorrowError : public BError
{
public:
        BBorrowError(const char* origin)
                :
                BError(origin)
        {
        }

        virtual const char* Message() const noexcept override { return "BBorrowError"; }
};


class BorrowAdmin
{
private:
        static constexpr uint8 kOwned = 0x1;
        static constexpr uint8 kBorrowed = 0x2;
        std::atomic<uint8> fState = kOwned;

protected:
        virtual ~BorrowAdmin() = default;

        virtual void Cleanup() noexcept {};
        virtual void ReleasePointer() noexcept {};

public:
        BorrowAdmin() noexcept {}


        void Borrow()
        {
                auto alreadyBorrowed = (fState.fetch_or(kBorrowed) & kBorrowed) == kBorrowed;
                if (alreadyBorrowed) {
                        throw BBorrowError(__PRETTY_FUNCTION__);
                }
        }


        void Return() noexcept
        {
                auto cleanup = (fState.fetch_and(~kBorrowed) & kOwned) != kOwned;
                if (cleanup)
                        this->Cleanup();
        }


        void Forfeit() noexcept
        {
                auto cleanup = (fState.fetch_and(~kOwned) & kBorrowed) != kBorrowed;
                if (cleanup)
                        this->Cleanup();
        }


        bool IsBorrowed() noexcept { return (fState.load() & kBorrowed) == kBorrowed; }


        void Release()
        {
                if ((fState.load() & kBorrowed) == kBorrowed)
                        throw BBorrowError(__PRETTY_FUNCTION__);
                this->ReleasePointer();
                this->Cleanup();
        }
};


template<typename T> class BorrowPointer : public BorrowAdmin
{
public:
        BorrowPointer(T* object) noexcept
                :
                fPtr(object)
        {
        }

        virtual ~BorrowPointer() { delete fPtr; }

protected:
        virtual void Cleanup() noexcept override { delete this; }

        virtual void ReleasePointer() noexcept override { fPtr = nullptr; }

private:
        T* fPtr;
};


template<typename T> class BExclusiveBorrow
{
        template<typename P> friend class BBorrow;

        T* fPtr = nullptr;
        BorrowAdmin* fAdminBlock = nullptr;

public:
        BExclusiveBorrow() noexcept {}


        BExclusiveBorrow(nullptr_t) noexcept {}


        BExclusiveBorrow(T* object)
        {
                fAdminBlock = new BorrowPointer<T>(object);
                fPtr = object;
        }

        ~BExclusiveBorrow()
        {
                if (fAdminBlock)
                        fAdminBlock->Forfeit();
        }

        BExclusiveBorrow(const BExclusiveBorrow&) = delete;


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


        BExclusiveBorrow(BExclusiveBorrow&& other) noexcept
        {
                if (fAdminBlock)
                        fAdminBlock->Forfeit();
                fAdminBlock = other.fAdminBlock;
                fPtr = other.fPtr;
                other.fAdminBlock = nullptr;
                other.fPtr = nullptr;
        }


        BExclusiveBorrow& operator=(BExclusiveBorrow&& other) noexcept
        {
                if (fAdminBlock)
                        fAdminBlock->Forfeit();
                fAdminBlock = other.fAdminBlock;
                fPtr = other.fPtr;
                other.fAdminBlock = nullptr;
                other.fPtr = nullptr;
                return *this;
        }


        bool HasValue() const noexcept { return bool(fPtr); }


        T& operator*() const
        {
                if (fAdminBlock && !fAdminBlock->IsBorrowed())
                        return *fPtr;
                throw BBorrowError(__PRETTY_FUNCTION__);
        }


        T* operator->() const
        {
                if (fAdminBlock && !fAdminBlock->IsBorrowed())
                        return fPtr;
                throw BBorrowError(__PRETTY_FUNCTION__);
        }


        std::unique_ptr<T> Release()
        {
                if (!fAdminBlock)
                        throw BBorrowError(__PRETTY_FUNCTION__);
                fAdminBlock->Release();
                auto retval = std::unique_ptr<T>(fPtr);
                fPtr = nullptr;
                fAdminBlock = nullptr;
                return retval;
        }
};


template<typename T> class BBorrow
{
        T* fPtr = nullptr;
        BorrowAdmin* fAdminBlock = nullptr;

public:
        BBorrow() noexcept {}


        BBorrow(nullptr_t) noexcept {}


        template<typename P>
        explicit BBorrow(BExclusiveBorrow<P>& owner)
                :
                fPtr(owner.fPtr),
                fAdminBlock(owner.fAdminBlock)
        {
                fAdminBlock->Borrow();
        }


        BBorrow(const BBorrow&) = delete;


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


        BBorrow(BBorrow&& other) noexcept
                :
                fPtr(other.fPtr),
                fAdminBlock(other.fAdminBlock)
        {
                other.fPtr = nullptr;
                other.fAdminBlock = nullptr;
        }


        BBorrow& operator=(BBorrow&& other) noexcept
        {
                if (fAdminBlock)
                        fAdminBlock->Return();

                fPtr = other.fPtr;
                fAdminBlock = other.fAdminBlock;
                other.fPtr = nullptr;
                other.fAdminBlock = nullptr;
                return *this;
        }


        ~BBorrow()
        {
                if (fAdminBlock)
                        fAdminBlock->Return();
        }


        bool HasValue() const noexcept { return bool(fPtr); }


        T& operator*() const
        {
                if (fPtr)
                        return *fPtr;
                throw BBorrowError(__PRETTY_FUNCTION__);
        }


        T* operator->() const
        {
                if (fPtr)
                        return fPtr;
                throw BBorrowError(__PRETTY_FUNCTION__);
        }


        void Return() noexcept
        {
                if (fAdminBlock)
                        fAdminBlock->Return();
                fAdminBlock = nullptr;
                fPtr = nullptr;
        }
};


template<class T, class... _Args>
BExclusiveBorrow<T>
make_exclusive_borrow(_Args&&... __args)
{
        auto guardedObject = std::make_unique<T>(std::forward<_Args>(__args)...);
        auto retval = BExclusiveBorrow<T>(guardedObject.get());
        guardedObject.release();
        return retval;
}


} // namespace Network

} // namespace BPrivate

#endif // _B_EXCLUSIVE_BORROW_H