#ifndef LOCK_SPEC_H
#define LOCK_SPEC_H

#include <vector>
#include <functional>
#include <atomic>
#include <stdexcept>
#include <spinLock.h>

namespace smo {

/**
 * @brief LockSet - Manages a collection of locks for acquisition/release
 */
class LockSet
{
public:
	/**
	 * @brief Constructor
	 * @param requiredLocks Vector of lock references that must be acquired
	 */
	LockSet(std::vector<std::reference_wrapper<SpinLock>> requiredLocks = {})
	: requiredLocks(std::move(requiredLocks)), allLocksAcquired(false)
	{}

	/**
	 * @brief Try to acquire all locks in order
	 * @return true if all locks were acquired, false otherwise
	 */
	bool tryAcquire()
	{
		if (allLocksAcquired)
		{
			throw std::runtime_error(
				std::string(__func__) +
				": LockSet::tryAcquire() called but allLocksAcquired is "
				"already true");
		}

		// Try to acquire all required locks
		int nAcquired = 0;
		for (auto& lockRef : requiredLocks)
		{
			if (!lockRef.get().tryAcquire()) { break; }
			nAcquired++;
		}

		if (nAcquired < static_cast<int>(requiredLocks.size()))
		{
			// Release any locks we managed to acquire
			for (int i = 0; i < nAcquired; i++) {
				requiredLocks[i].get().release();
			}
			allLocksAcquired = false;
			return false;
		}

		allLocksAcquired = true;
		return true;
	}

	/**
	 * @brief Release all locks
	 */
	void release()
	{
		if (!allLocksAcquired)
		{
			throw std::runtime_error(
				std::string(__func__) +
				": LockSet::release() called but allLocksAcquired is false");
		}

		for (auto& lockRef : requiredLocks) {
			lockRef.get().release();
		}

		allLocksAcquired = false;
	}

private:
    std::vector<std::reference_wrapper<SpinLock>> requiredLocks;
    std::atomic<bool> allLocksAcquired;
};

} // namespace smo

#endif // LOCK_SPEC_H