#ifndef FZ_BORROW_HPP
#define FZ_BORROW_HPP

#include "rundown_protection.hpp"

namespace fz {

// Forward declaration
template<typename T>
class scoped_borrow;

/**
 * Base class that makes a type "borrowable"
 */
class borrowable {
public:
	borrowable() = default;

	// Non-copyable, non-movable (borrowed objects shouldn't be moved)
	borrowable(const borrowable&) = delete;
	borrowable& operator=(const borrowable&) = delete;
	borrowable(borrowable&&) = delete;
	borrowable& operator=(borrowable&&) = delete;

	bool is_borrowable() const noexcept
	{
		return rundown_protection_.acquire_enabled();
	}

	void stop_borrowing() noexcept {
		rundown_protection_.disable_acquire();
	}

	void wait_until_no_active_borrows() noexcept {
		rundown_protection_.wait_until_all_acquires_released();
	}

private:
	template <typename T>
	friend class scoped_borrow;

	rundown_protection rundown_protection_;
};

/**
 * RAII guard for borrowing an object. Similar to fz::scoped_lock.
 * Automatically releases the borrow when destroyed.
 */
template<typename T>
class scoped_borrow {
public:
	/**
	 * Default constructor - creates an invalid scoped_borrow.
	 */
	scoped_borrow() noexcept
		: ptr_(nullptr)
		, b_(nullptr)
	{}

	/**
	 * Construct by attempting to borrow from a raw pointer.
	 * The scoped_borrow will be invalid if the borrow fails.
	 */
	template <typename U = T, std::enable_if_t<std::is_base_of_v<borrowable, U>>* = nullptr>
	explicit scoped_borrow(T *ptr) noexcept
		: scoped_borrow()
	{
		if (ptr && ptr->borrowable::rundown_protection_.acquire()) {
			ptr_ = ptr;
			b_ = ptr;
		}
	}

	/**
	 * Construct by attempting to borrow from a unique_ptr.
	 */
	explicit scoped_borrow(const std::unique_ptr<T> &ptr) noexcept
		: scoped_borrow(ptr.get())
	{}

	/**
	 * Construct by attempting to borrow from a shared_ptr.
	 */
	explicit scoped_borrow(const std::shared_ptr<T> &ptr) noexcept
		: scoped_borrow(ptr.get())
	{}

	/**
	 * Aliasing constructor
	 * The scoped_borrow will be invalid if the borrow fails.
	 */
	explicit scoped_borrow(T *ptr, borrowable *b) noexcept
		: scoped_borrow()

	{
		if (ptr && b && b->rundown_protection_.acquire()) {
			ptr_ = ptr;
			b_ = b;
		}
	}

	~scoped_borrow()
	{
		reset();
	}

	// Non-copyable
	scoped_borrow(const scoped_borrow &) = delete;
	scoped_borrow& operator=(const scoped_borrow &) = delete;

	friend void swap(scoped_borrow &lhs, scoped_borrow &rhs) noexcept
	{
		using std::swap;

		swap(lhs.ptr_, rhs.ptr_);
		swap(lhs.b_, rhs.b_);
	}

	scoped_borrow(scoped_borrow &&other) noexcept
		: ptr_(other.ptr_)
		, b_(other.b_)
	{
		other.ptr_ = nullptr;
		other.b_ = nullptr;
	}

	scoped_borrow& operator=(scoped_borrow &&other) noexcept
	{
		scoped_borrow moved(std::move(other));

		swap(*this, moved);

		return *this;
	}

	explicit operator bool() const noexcept
	{
		return ptr_ != nullptr;
	}

	T* operator->() const noexcept
	{
		return ptr_;
	}

	T &operator*() const noexcept
	{
		return *ptr_;
	}

	void reset() noexcept
	{
		if (b_) {
			b_->rundown_protection_.release();
			ptr_ = nullptr;
			b_ = nullptr;
		}
	}

private:
	template <typename A, typename B>
	friend scoped_borrow<A> static_scoped_borrow_cast(scoped_borrow<B> &&b);

	template <typename A, typename B>
	friend scoped_borrow<A> dynamic_scoped_borrow_cast(scoped_borrow<B> &&b);

	template <typename A, typename B>
	friend scoped_borrow<A> reinterpret_scoped_borrow_cast(scoped_borrow<B> &&b);

	T* ptr_;
	borrowable *b_;
};

template <typename A, typename B>
inline scoped_borrow<A> static_scoped_borrow_cast(scoped_borrow<B> &&b)
{
	scoped_borrow<A> ret;

	ret.ptr_ = static_cast<A*>(b.ptr_);
	ret.b_ = b.b_;

	b.ptr_ = nullptr;
	b.b_ = nullptr;

	return ret;
}

template <typename A, typename B>
inline scoped_borrow<A> dynamic_scoped_borrow_cast(scoped_borrow<B> &&b)
{
	if (A *casted = dynamic_cast<A*>(b.ptr_)) {
		scoped_borrow<A> ret;

		ret.ptr_ = casted;
		ret.b_ = b.b_;

		b.ptr_ = nullptr;
		b.b_ = nullptr;

		return ret;
	}

	return {};
}

template <typename A, typename B>
inline scoped_borrow<A> reinterpret_scoped_borrow_cast(scoped_borrow<B> &&b)
{
	scoped_borrow<A> ret;

	ret.ptr_ = reinterpret_cast<A*>(b.ptr_);
	ret.b_ = b.b_;

	b.ptr_ = nullptr;
	b.b_ = nullptr;

	return ret;
}

}

#endif // FZ_BORROW_HPP
