#include <user/stagingBuffer.h>
#include <cassert>
#include <unistd.h>
#include <cstdint>
#include <stdexcept>
#include <sys/mman.h>
#include <vector>
#include <liburing.h>

#include <user/frameAssemblyDesc.h>

namespace smo {
namespace stim_buff {

static const char* pinningMechanismToString(
	StagingBuffer::PinningMechanism mechanism)
{
	switch (mechanism)
	{
		case StagingBuffer::PinningMechanism::NONE:
			return "NONE";
		case StagingBuffer::PinningMechanism::MLOCK:
			return "MLOCK";
		case StagingBuffer::PinningMechanism::IO_URING:
			return "IO_URING";
	}

	return "Unknown";
}

// Static defaults for io_uring
const StagingBuffer::IOEngineConstraints
StagingBuffer::IOEngineConstraints::ioUringConstraints(
	// slotStartAlignmentByteVal (page alignment for DMA)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE)),
	// slotPadToNBytes (MTU 1500 - UDP/IP header 28)
	1472,
	// frameStartAlignmentByteVal (page alignment for DMA)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE)),
	// framePadToNBytes (MTU 1500 - UDP/IP header 28)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE))
);

// Static defaults for OpenCL input
const StagingBuffer::IOEngineConstraints
StagingBuffer::IOEngineConstraints::openClInputConstraints(
	// slotStartAlignmentByteVal (page alignment)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE)),
	// slotPadToNBytes (XYZI point size)
	16,
	// frameStartAlignmentByteVal (page alignment)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE)),
	// framePadToNBytes (pointer size)
	static_cast<size_t>(sysconf(_SC_PAGE_SIZE))
);

// Helper function to calculate maximum alignment needed for first slot
// (must satisfy both frame and slot alignment)
static size_t calculateMaxAlignment(
	size_t frameStartAlignmentByteVal,
	size_t slotStartAlignmentByteVal)
{
	if (frameStartAlignmentByteVal >= slotStartAlignmentByteVal)
	{
		if (frameStartAlignmentByteVal % slotStartAlignmentByteVal == 0)
			{ return frameStartAlignmentByteVal; }
		else
		{
			// Need LCM, but for simplicity use the larger alignment
			// In practice, alignments are usually powers of 2, so this should work
			return std::max(
				frameStartAlignmentByteVal, slotStartAlignmentByteVal);
		}
	}
	else
	{
		if (slotStartAlignmentByteVal % frameStartAlignmentByteVal == 0)
			{ return slotStartAlignmentByteVal; }
		else
		{
			return std::max(
				frameStartAlignmentByteVal, slotStartAlignmentByteVal);
		}
	}
}

void StagingBuffer::computeSlotStrideAndBufferSize()
{
	// Slot stride is the maximum of alignment and padding, rounded up to a multiple of alignment
	size_t minSlotStride = std::max(
		inputConstraints.slotStartAlignmentByteVal,
		inputConstraints.slotPadToNBytes);
	slotStrideNBytes = ((minSlotStride + inputConstraints.slotStartAlignmentByteVal - 1)
		/ inputConstraints.slotStartAlignmentByteVal)
		* inputConstraints.slotStartAlignmentByteVal;

	// Calculate maximum alignment needed for first slot (must satisfy both frame and slot alignment)
	size_t maxAlignment = calculateMaxAlignment(
		inputConstraints.frameStartAlignmentByteVal,
		inputConstraints.slotStartAlignmentByteVal);

	// Calculate minimum buffer size
	size_t minBufferSize = std::max(
		inputConstraints.framePadToNBytes,
		inputConstraints.slotPadToNBytes);

	// Calculate total size needed for nSlots slots
	size_t slotAreaSize = nSlots * slotStrideNBytes;

	// Add padding space at buffer start for alignment offset (worst case: max alignment - 1)
	size_t alignmentPadding = maxAlignment - 1;

	// Total size needed: alignment padding + slot area, then ensure minimum is met
	size_t rawSize = alignmentPadding + slotAreaSize;
	if (rawSize < minBufferSize)
		{ rawSize = minBufferSize; }

	// Align up to the maximum alignment to ensure we can always find a valid offset
	bufferNBytes = ((rawSize + maxAlignment - 1) / maxAlignment) * maxAlignment;
}

// Static member function to calculate offset and validate invariants
size_t StagingBuffer::calculateFirstSlotOffsetAndValidate(
	uint8_t* buffer,
	size_t bufferNBytes,
	size_t nSlots,
	size_t slotStrideNBytes,
	const StagingBuffer::IOEngineConstraints& inputConstraints)
{
	// Calculate maximum alignment needed for first slot
	size_t maxAlignment = calculateMaxAlignment(
		inputConstraints.frameStartAlignmentByteVal,
		inputConstraints.slotStartAlignmentByteVal);

	// Calculate offset to align first slot to both frame and slot alignment
	uintptr_t bufferAddr = reinterpret_cast<uintptr_t>(buffer);
	uintptr_t alignedAddr = ((bufferAddr + maxAlignment - 1) / maxAlignment)
		* maxAlignment;
	size_t firstSlotOffsetNBytes = alignedAddr - bufferAddr;

	// Validate invariants with exceptions
	uint8_t* firstSlotAddr = buffer + firstSlotOffsetNBytes;
	if (
		reinterpret_cast<uintptr_t>(firstSlotAddr)
		% inputConstraints.frameStartAlignmentByteVal != 0)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": StagingBuffer: first slot address not aligned to "
			+ std::to_string(inputConstraints.frameStartAlignmentByteVal));
	}

	if (
		reinterpret_cast<uintptr_t>(firstSlotAddr)
		% inputConstraints.slotStartAlignmentByteVal != 0)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": StagingBuffer: first slot address not aligned to "
			+ std::to_string(inputConstraints.slotStartAlignmentByteVal));
	}

	size_t minBufferSize = std::max(
		inputConstraints.framePadToNBytes,
		inputConstraints.slotPadToNBytes);
	if (bufferNBytes < minBufferSize)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": StagingBuffer: buffer size less than minimum required (max of "
			+ std::to_string(inputConstraints.framePadToNBytes)
			+ " and "
			+ std::to_string(inputConstraints.slotPadToNBytes)
			+ ")");
	}

	if (firstSlotOffsetNBytes + nSlots * slotStrideNBytes
		> bufferNBytes)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": StagingBuffer: buffer size insufficient to hold "
			+ std::to_string(nSlots)
			+ " slots with proper alignment and padding");
	}

	return firstSlotOffsetNBytes;
}

StagingBuffer::StagingBuffer(
	const IOEngineConstraints& inputEngineConstraints_,
	const IOEngineConstraints& /*outputEngineConstraints*/,
	size_t nSlots)
: buffer(nullptr, MmapDeleter(0)),
nSlots(nSlots),
inputConstraints(inputEngineConstraints_)
{
	if (nSlots == 0)
	{
		throw std::invalid_argument(std::string(__func__)
			+ ": StagingBuffer: nSlots must be > 0");
	}

	computeSlotStrideAndBufferSize();

	/* Allocate buffer using mmap() for io_uring registration
	 * MAP_ANONYMOUS | MAP_PRIVATE creates anonymous, non-file-backed memory
	 */
	void* mmapped = mmap(
		nullptr, bufferNBytes,
		PROT_READ | PROT_WRITE,
		MAP_ANONYMOUS | MAP_PRIVATE,
		-1, 0);

	if (mmapped == MAP_FAILED)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": StagingBuffer: mmap() failed");
	}

	buffer = std::unique_ptr<uint8_t, MmapDeleter>(
		static_cast<uint8_t*>(mmapped), MmapDeleter(bufferNBytes));
	currentNBytes.store(0);

	// Calculate offset and validate invariants (helper function in .cpp)
	firstSlotOffsetNBytes = StagingBuffer::calculateFirstSlotOffsetAndValidate(
		buffer.get(), bufferNBytes, nSlots,
		slotStrideNBytes, inputConstraints);

	// Build FrameAssemblyDesc once
	std::vector<FrameAssemblyDesc::SlotDesc> slots;
	slots.reserve(nSlots);
	uint8_t *frameBase = buffer.get() + firstSlotOffsetNBytes;
	for (size_t i = 0; i < nSlots; ++i)
	{
		size_t off = i * slotStrideNBytes;
		FrameAssemblyDesc::SlotDesc s{
			off, frameBase + off, inputConstraints.slotPadToNBytes};

		slots.push_back(s);
	}

	frameDesc = std::make_shared<FrameAssemblyDesc>(
		nSlots, inputConstraints.slotPadToNBytes, bufferNBytes,
		std::move(slots));
}

StagingBuffer::~StagingBuffer()
{
	assert(!currentlyPinned);
}

StagingBuffer::Pinner::Pinner(StagingBuffer& parent_)
: parent(parent_)
{}

void StagingBuffer::assertUnpinnedAndMarkPinned(PinningMechanism mechanism)
{
	if (currentlyPinned)
	{
		throw std::runtime_error(
			std::string(__func__) + ": StagingBuffer already pinned with "
			+ pinningMechanismToString(currentPinningMechanism));
	}

	currentlyPinned = true;
	currentPinningMechanism = mechanism;
}

std::unique_ptr<StagingBuffer::MlockPinner> StagingBuffer::makeMlockPinner()
{
	return std::make_unique<MlockPinner>(*this);
}

std::unique_ptr<StagingBuffer::IoUringPinner> StagingBuffer::makeIoUringPinner(
	struct io_uring* ring)
{
	return std::make_unique<IoUringPinner>(*this, ring);
}

StagingBuffer::MlockPinner::MlockPinner(StagingBuffer& parent_)
: Pinner(parent_)
{
	if (!parent.buffer || parent.bufferNBytes == 0)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": Cannot mlock an uninitialized StagingBuffer");
	}

	parent.assertUnpinnedAndMarkPinned(PinningMechanism::MLOCK);

	if (mlock(parent.buffer.get(), parent.bufferNBytes) != 0)
	{
		parent.currentlyPinned = false;
		parent.currentPinningMechanism = PinningMechanism::NONE;
		throw std::runtime_error(std::string(__func__)
			+ ": mlock() failed");
	}
}

StagingBuffer::MlockPinner::~MlockPinner()
{
	assert(parent.currentlyPinned);
	assert(parent.currentPinningMechanism == PinningMechanism::MLOCK);

	int ret = munlock(parent.buffer.get(), parent.bufferNBytes);
	assert(ret == 0);
	(void)ret;

	parent.currentlyPinned = false;
	parent.currentPinningMechanism = PinningMechanism::NONE;
}

StagingBuffer::IoUringPinner::IoUringPinner(
	StagingBuffer& parent_, struct io_uring* ring_)
: Pinner(parent_), ring(ring_)
{
	if (!ring)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": io_uring ring pointer is null");
	}

	if (!parent.buffer || parent.bufferNBytes == 0)
	{
		throw std::runtime_error(std::string(__func__)
			+ ": Cannot register an uninitialized StagingBuffer");
	}

	parent.assertUnpinnedAndMarkPinned(PinningMechanism::IO_URING);

	struct iovec iov = parent.getIoUringRegisterIoVec();
	int ret = io_uring_register_buffers(ring, &iov, 1);
	if (ret < 0)
	{
		parent.currentlyPinned = false;
		parent.currentPinningMechanism = PinningMechanism::NONE;
		throw std::runtime_error(std::string(__func__)
			+ ": io_uring_register_buffers failed");
	}
}

StagingBuffer::IoUringPinner::~IoUringPinner()
{
	assert(parent.currentlyPinned);
	assert(parent.currentPinningMechanism == PinningMechanism::IO_URING);

	int ret = io_uring_unregister_buffers(ring);
	assert(ret == 0);
	(void)ret;

	parent.currentlyPinned = false;
	parent.currentPinningMechanism = PinningMechanism::NONE;
}

} // namespace stim_buff
} // namespace smo
