#include <gtest/gtest.h>
#include <user/stagingBuffer.h>
#include <user/frameAssemblyDesc.h>
#include <unistd.h>
#include <cstdint>
#include <algorithm>
#include <string>

namespace smo {
namespace stim_buff {

// Helper function to create test constraints
static StagingBuffer::IOEngineConstraints createTestConstraints(
	size_t slotStartAlignment,
	size_t slotPadToNBytes,
	size_t frameStartAlignment,
	size_t framePadToNBytes)
{
	return StagingBuffer::IOEngineConstraints(
		slotStartAlignment,
		slotPadToNBytes,
		frameStartAlignment,
		framePadToNBytes);
}

// Helper function to verify alignment
static void verifyAlignment(
	uint8_t* addr,
	size_t alignment,
	const std::string& description)
{
	uintptr_t addrValue = reinterpret_cast<uintptr_t>(addr);
	EXPECT_EQ(addrValue % alignment, 0u)
		<< description << ": Address " << (void*)addr
		<< " is not aligned to " << alignment;
}

class StagingBufferTest : public ::testing::Test {
protected:
	void SetUp() override {
		pageSize = static_cast<size_t>(sysconf(_SC_PAGE_SIZE));
	}

	void TearDown() override {
	}

	size_t pageSize;
};

// Helper function to verify all slots are page-aligned and have correct stride
static void verifyAllSlotsPageAligned(
	const StagingBuffer& buffer,
	size_t expectedSlotStride,
	size_t pageSize)
{
	auto frameDesc = static_cast<std::shared_ptr<FrameAssemblyDesc>>(buffer);
	ASSERT_NE(frameDesc, nullptr);
	ASSERT_GT(frameDesc->slots.size(), 0u);

	// Verify slot stride is a multiple of page size
	EXPECT_EQ(expectedSlotStride % pageSize, 0u)
		<< "Slot stride " << expectedSlotStride
		<< " should be a multiple of page size " << pageSize;

	// Verify first slot is page-aligned
	verifyAlignment(frameDesc->slots[0].vaddr, pageSize, "First slot");

	// Verify all subsequent slots are page-aligned and have correct stride
	for (size_t i = 1; i < frameDesc->slots.size(); ++i) {
		verifyAlignment(frameDesc->slots[i].vaddr, pageSize,
			"Slot " + std::to_string(i));

		// Verify actual stride matches expected stride
		size_t actualStride = reinterpret_cast<uintptr_t>(frameDesc->slots[i].vaddr) -
			reinterpret_cast<uintptr_t>(frameDesc->slots[i-1].vaddr);
		EXPECT_EQ(actualStride, expectedSlotStride)
			<< "Slot " << i << " stride mismatch: expected " << expectedSlotStride
			<< ", got " << actualStride;
	}
}

// Test 1: Small slot stride (< page size) - should round up to page size
TEST_F(StagingBufferTest, SmallSlotStrideRoundsUpToPageSize) {
	size_t nSlots = 10;
	size_t smallSlotPad = 256; // Much smaller than typical page size (4096)
	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment (page size)
		smallSlotPad,            // slotPadToNBytes (small)
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Slot stride should be rounded up to page size
	size_t expectedSlotStride = pageSize;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Small slot pad should round up to page size";

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);
}

// Test 2: Slot stride equal to page size - should remain page size
TEST_F(StagingBufferTest, SlotStrideEqualToPageSize) {
	size_t nSlots = 20;
	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		pageSize,                // slotPadToNBytes (equal to page size)
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	size_t expectedSlotStride = pageSize;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Slot stride equal to page size should remain unchanged";

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);
}

// Test 3: Slot stride slightly larger than page size (e.g., 336) - should round up
TEST_F(StagingBufferTest, SlotStrideSlightlyLargerThanPageSize) {
	size_t nSlots = 50;
	size_t slotPadToNBytes = 336; // Slightly larger than page size (4096)
	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		slotPadToNBytes,         // slotPadToNBytes
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Should round up to next multiple of page size (4096)
	size_t expectedSlotStride = pageSize;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Slot stride 336 should round up to page size " << pageSize;

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);
}

// Test 4: Slot stride much larger than page size (e.g., 32256) - should round up to 32768
TEST_F(StagingBufferTest, SlotStrideMuchLargerThanPageSize) {
	size_t nSlots = 100;
	size_t slotPadToNBytes = 32256; // Much larger than page size
	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		slotPadToNBytes,         // slotPadToNBytes
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Should round up to next multiple of page size
	// 32256 / 4096 = 7.875, so rounds up to 8 pages = 32768
	size_t expectedSlotStride = ((slotPadToNBytes + pageSize - 1) / pageSize) * pageSize;
	EXPECT_EQ(expectedSlotStride, 32768u)
		<< "32256 should round up to 32768 (8 pages)";
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Slot stride should be rounded up to " << expectedSlotStride;

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);
}

// Test 5: Slot stride already a multiple of page size - should remain unchanged
TEST_F(StagingBufferTest, SlotStrideAlreadyMultipleOfPageSize) {
	size_t nSlots = 30;
	size_t slotPadToNBytes = pageSize * 3; // Already a multiple (e.g., 12288)
	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		slotPadToNBytes,         // slotPadToNBytes (already multiple)
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Should remain unchanged
	size_t expectedSlotStride = slotPadToNBytes;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Slot stride already a multiple of page size should remain unchanged";

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);
}

// Test 6: Multiple different slot stride sizes in sequence
TEST_F(StagingBufferTest, MultipleDifferentSlotStrideSizes) {
	// Test with various slot pad sizes
	struct TestCase {
		size_t slotPadToNBytes;
		size_t expectedRoundedStride;
	};

	std::vector<TestCase> testCases = {
		{256, pageSize},                    // Small: rounds to 1 page
		{pageSize, pageSize},               // Equal: stays 1 page
		{pageSize + 1, pageSize * 2},      // Slightly larger: rounds to 2 pages
		{32256, 32768},                     // Much larger: rounds to 8 pages
		{pageSize * 5, pageSize * 5},       // Already multiple: stays 5 pages
		{pageSize * 10 + 100, pageSize * 11}, // Large with remainder: rounds to 11 pages
	};

	for (const auto& testCase : testCases) {
		size_t nSlots = 10;
		auto constraints = createTestConstraints(
			pageSize,
			testCase.slotPadToNBytes,
			pageSize,
			pageSize);

		StagingBuffer buffer(constraints, constraints, nSlots);

		EXPECT_EQ(buffer.slotStrideNBytes, testCase.expectedRoundedStride)
			<< "Slot pad " << testCase.slotPadToNBytes
			<< " should round to " << testCase.expectedRoundedStride;

		verifyAllSlotsPageAligned(buffer, testCase.expectedRoundedStride, pageSize);
	}
}

// Test 7: Real-world scenario - PcloudIntensityStimulusBuffer (32256 bytes)
TEST_F(StagingBufferTest, RealWorldPcloudIntensityScenario) {
	size_t nSlots = 909; // histbuffMs=30000 / CONFIG_STIMBUFF_FRAME_PERIOD_MS=33
	size_t nDgramsPerFrame = 84;
	size_t nPointsPerDgram = 96;
	size_t slotPadToNBytes = nDgramsPerFrame * nPointsPerDgram * sizeof(float) * 1; // 32256

	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		slotPadToNBytes,         // slotPadToNBytes
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Should round up to 32768 (8 pages)
	size_t expectedSlotStride = 32768;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Real-world 32256-byte slots should round to 32768";

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);

	// Verify FrameAssemblyDesc
	auto frameDesc = static_cast<std::shared_ptr<FrameAssemblyDesc>>(buffer);
	ASSERT_NE(frameDesc, nullptr);
	EXPECT_EQ(frameDesc->numSlots, nSlots);
	EXPECT_EQ(frameDesc->slotSizeBytes, slotPadToNBytes);
}

// Test 8: Real-world scenario - PcloudAmbienceStimulusBuffer (336 bytes)
TEST_F(StagingBufferTest, RealWorldPcloudAmbienceScenario) {
	size_t nSlots = 909;
	size_t nDgramsPerFrame = 84;
	size_t slotPadToNBytes = nDgramsPerFrame * sizeof(float); // 336

	auto constraints = createTestConstraints(
		pageSize,                // slotStartAlignment
		slotPadToNBytes,         // slotPadToNBytes
		pageSize,                // frameStartAlignment
		pageSize);               // framePadToNBytes

	StagingBuffer buffer(constraints, constraints, nSlots);

	// Should round up to page size
	size_t expectedSlotStride = pageSize;
	EXPECT_EQ(buffer.slotStrideNBytes, expectedSlotStride)
		<< "Real-world 336-byte slots should round to page size";

	verifyAllSlotsPageAligned(buffer, expectedSlotStride, pageSize);

	// Verify FrameAssemblyDesc
	auto frameDesc = static_cast<std::shared_ptr<FrameAssemblyDesc>>(buffer);
	ASSERT_NE(frameDesc, nullptr);
	EXPECT_EQ(frameDesc->numSlots, nSlots);
	EXPECT_EQ(frameDesc->slotSizeBytes, slotPadToNBytes);
}

} // namespace stim_buff
} // namespace smo
