#ifndef STAGINGBUFFER_H
#define STAGINGBUFFER_H

#include <memory>
#include <cstdint>
#include <atomic>
#include <string>
#include <sstream>
#include <sys/mman.h>
#include <sys/uio.h>

namespace smo {
namespace stim_buff {

// Forward declaration
class FrameAssemblyDesc;

/**
 * StagingBuffer manages a large buffer to guide io_uring in assembling some
 * number of Livox Avia pcloud UDP dgrams into a single stim frame.
 *
 * The buffer operates in a cycle:
 * 1. io_uring assembles UDP dgrams into the buffer until it's full
 * 2. Buffer is handed off to the stimbuff layer to be appended to the stimbuff.
 * 3. When the stimbuff layer has appended the current assembled frame, the
 *		assembly buffer is reset and cycle repeats.
 */
class StagingBuffer
{
public:
	class IOEngineConstraints
	{
	public:
		IOEngineConstraints(
			size_t slotStartAlignmentByteVal_,
			size_t slotPadToNBytes_,
			size_t frameStartAlignmentByteVal_,
			size_t framePadToNBytes_)
		: slotStartAlignmentByteVal(slotStartAlignmentByteVal_),
		slotPadToNBytes(slotPadToNBytes_),
		frameStartAlignmentByteVal(frameStartAlignmentByteVal_),
		framePadToNBytes(framePadToNBytes_)
		{}

		~IOEngineConstraints() = default;

		size_t slotStartAlignmentByteVal, slotPadToNBytes,
			frameStartAlignmentByteVal, framePadToNBytes;

		// Static defaults for io_uring and OpenCL
		static const IOEngineConstraints ioUringConstraints;
		static const IOEngineConstraints openClInputConstraints;

		inline std::string stringify() const
		{
			std::ostringstream oss;
			oss << "IOEngineConstraints{"
				<< "slotStartAlignmentByteVal=" << slotStartAlignmentByteVal
				<< ", slotPadToNBytes=" << slotPadToNBytes
				<< ", frameStartAlignmentByteVal=" << frameStartAlignmentByteVal
				<< ", framePadToNBytes=" << framePadToNBytes
				<< "}";
			return oss.str();
		}
	};

public:
	/**	EXPLANATION:
	 * We use the input and output engine constraints to determine the total
	 * amount of memory required internally to assemble a single frame with
	 * the given number of points per frame.
	 */
	explicit StagingBuffer(
		const IOEngineConstraints& inputEngineConstraints,
		const IOEngineConstraints& outputEngineConstraints,
		size_t nDgramsPerFrame);
	~StagingBuffer() = default;

	// Non-copyable, movable
	StagingBuffer(const StagingBuffer&) = delete;
	StagingBuffer& operator=(const StagingBuffer&) = delete;
	StagingBuffer(StagingBuffer&&) = default;
	StagingBuffer& operator=(StagingBuffer&&) = default;

public:
	/**	EXPLANATION:
	 * Returns an input-engine-agnostic descriptor describing per-frame packet
	 * slot layout. Different input engines should be able to convert this into
	 * engine-specific metadata. E.g: io_uring's SQE descriptor.
	 */
	operator std::shared_ptr<FrameAssemblyDesc>() const { return frameDesc; }
	// operator OpenClSharedBufferDescriptor() const;

	bool isAssembling() const { return assemblingFlag.load(); }
	void startAssembly() { assemblingFlag.store(true); }
	void stopAssembly() { assemblingFlag.store(false); }

	/**	EXPLANATION:
	 * Returns an iovec for io_uring registration.
	 * The buffer is mmap()-allocated and suitable for IORING_REGISTER_BUFFERS.
	 */
	struct iovec getIoUringRegisterIoVec() const
	{
		struct iovec iov;
		iov.iov_base = buffer.get();
		iov.iov_len = bufferNBytes;
		return iov;
	}

	/**	EXPLANATION:
	 * Returns an iovec for OpenCL engine buffer access.
	 * The buffer is mmap()-allocated and suitable for CL_MEM_USE_HOST_PTR.
	 */
	struct iovec getClEngineIovec() const
	{
		struct iovec iov;
		iov.iov_base = buffer.get();
		iov.iov_len = bufferNBytes;
		return iov;
	}

	inline std::string stringify() const
	{
		std::ostringstream oss;
		oss << "StagingBuffer{"
			<< "nDgramsPerFrame=" << nDgramsPerFrame
			<< ", bufferNBytes=" << bufferNBytes
			<< ", slotStrideNBytes=" << slotStrideNBytes
			<< ", constraints=" << inputConstraints.stringify()
			<< "}";
		return oss.str();
	}

private:
	void computeSlotStrideAndBufferSize();
	static size_t calculateFirstSlotOffsetAndValidate(
		uint8_t* buffer,
		size_t bufferNBytes,
		size_t nDgramsPerFrame,
		size_t slotStrideNBytes,
		const IOEngineConstraints& inputConstraints);

	// Custom deleter for mmap-allocated buffer
	struct MmapDeleter
	{
		size_t size;
		MmapDeleter(size_t s) : size(s) {}
		
		void operator()(uint8_t* ptr) const
		{
			if (ptr != nullptr && size > 0)
			{
				munlock(ptr, size);
				munmap(ptr, size);
			}
		}
	};

	// Buffer data - mmap-allocated for io_uring registration
	// Using unique_ptr<uint8_t, MmapDeleter> instead of array syntax
	// since we have a custom deleter that knows the size
	std::unique_ptr<uint8_t, MmapDeleter> buffer;
	size_t bufferNBytes;

	// Layout/invariants
	size_t nDgramsPerFrame;

public:
	size_t slotStrideNBytes;
	size_t firstSlotOffsetNBytes;  // offset from buffer start to first slot

private:
	IOEngineConstraints inputConstraints;

	// Descriptor (computed once; reused across frames)
	mutable std::shared_ptr<FrameAssemblyDesc> frameDesc;

	// Current state
	std::atomic<size_t> currentNBytes;
	std::atomic<bool> assemblingFlag;
};

} // namespace stim_buff
} // namespace smo

#endif // STAGINGBUFFER_H