/*
 * Copyright (C) 2023-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "shared/source/utilities/heap_allocator.h"
#include "shared/test/common/fixtures/device_fixture.h"
#include "shared/test/common/mocks/mock_device.h"
#include "shared/test/common/mocks/mock_product_helper.h"
#include "shared/test/common/test_macros/test.h"

#include "gtest/gtest.h"

using namespace NEO;

class MockISAPool : public ISAPool {
  public:
    using ISAPool::chunkAllocator;
};

class MockISAPoolAllocator : public ISAPoolAllocator {
  public:
    using ISAPoolAllocator::bufferPools;

    MockISAPoolAllocator(Device *device) : ISAPoolAllocator(device) {}
};

using IsaPoolAllocatorTest = Test<DeviceFixture>;

void verifySharedIsaAllocation(const SharedIsaAllocation *sharedAllocation, size_t expectedOffset, size_t expectedSize) {
    ASSERT_NE(nullptr, sharedAllocation);
    EXPECT_NE(nullptr, sharedAllocation->getGraphicsAllocation());
    EXPECT_EQ(expectedOffset, sharedAllocation->getOffset());
    EXPECT_EQ(expectedSize, sharedAllocation->getSize());
}

TEST_F(IsaPoolAllocatorTest, givenIsaPoolAllocatorWhenNoAllocationsThenCreateNewAllocationPerIsaType) {
    auto &isaAllocator = pDevice->getIsaPoolAllocator();
    constexpr size_t requestAllocationSize = MemoryConstants::pageSize;

    // Allocate for user and builtin and verify that these are different allocations
    auto userAllocation = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
    verifySharedIsaAllocation(userAllocation, 0ul, requestAllocationSize);
    EXPECT_EQ(AllocationType::kernelIsa, userAllocation->getGraphicsAllocation()->getAllocationType());

    auto builtinAllocation = isaAllocator.requestGraphicsAllocationForIsa(true, requestAllocationSize);
    verifySharedIsaAllocation(builtinAllocation, 0ul, requestAllocationSize);
    EXPECT_EQ(AllocationType::kernelIsaInternal, builtinAllocation->getGraphicsAllocation()->getAllocationType());

    EXPECT_NE(userAllocation->obtainSharedAllocationLock().mutex(), builtinAllocation->obtainSharedAllocationLock().mutex());
    EXPECT_NE(userAllocation->getGraphicsAllocation(), builtinAllocation->getGraphicsAllocation());
    isaAllocator.freeSharedIsaAllocation(userAllocation);
    isaAllocator.freeSharedIsaAllocation(builtinAllocation);
}

TEST_F(IsaPoolAllocatorTest, givenIsaPoolAllocatorWhenAllocationsExistThenReuseAllocation) {
    auto &isaAllocator = pDevice->getIsaPoolAllocator();

    constexpr size_t requestAllocationSize = MemoryConstants::pageSize;

    auto allocation = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
    verifySharedIsaAllocation(allocation, 0ul, requestAllocationSize);

    auto allocationSize = allocation->getGraphicsAllocation()->getUnderlyingBufferSize();
    auto numOfSharedAllocations = allocationSize / requestAllocationSize;

    // Perform requests until allocation is full
    for (auto i = 1u; i < numOfSharedAllocations; i++) {
        auto tempSharedAllocation = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
        verifySharedIsaAllocation(tempSharedAllocation, requestAllocationSize * (i), requestAllocationSize);
        EXPECT_EQ(allocation->getGraphicsAllocation(), tempSharedAllocation->getGraphicsAllocation());
        isaAllocator.freeSharedIsaAllocation(tempSharedAllocation);
    }

    // Verify that draining freed chunks is correct and allocation can be reused
    auto newAllocation = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
    verifySharedIsaAllocation(newAllocation, requestAllocationSize, requestAllocationSize);
    EXPECT_EQ(allocation->getGraphicsAllocation(), newAllocation->getGraphicsAllocation());

    isaAllocator.freeSharedIsaAllocation(newAllocation);
    isaAllocator.freeSharedIsaAllocation(allocation);
}

TEST_F(IsaPoolAllocatorTest, givenIsaPoolAllocatorWhenNoPoolAvailableThenCreateNewPool) {
    auto &isaAllocator = pDevice->getIsaPoolAllocator();

    constexpr size_t sharedIsaAllocationSize = MemoryConstants::pageSize2M * 2;
    constexpr size_t requestAllocationSize = sharedIsaAllocationSize / 2 + MemoryConstants::pageSize;

    constexpr size_t numOfAllocations = 5;
    SharedIsaAllocation *sharedAllocations[numOfAllocations];
    for (auto i = 0u; i < numOfAllocations; i++) {
        sharedAllocations[i] = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
        verifySharedIsaAllocation(sharedAllocations[i], 0, requestAllocationSize);
    }

    auto ga = sharedAllocations[0]->getGraphicsAllocation();
    for (auto i = 1u; i < numOfAllocations; i++) {
        EXPECT_NE(ga, sharedAllocations[i]->getGraphicsAllocation());
        isaAllocator.freeSharedIsaAllocation(sharedAllocations[i]);
    }
    isaAllocator.freeSharedIsaAllocation(sharedAllocations[0]);
}

TEST_F(IsaPoolAllocatorTest, givenIsaPoolAllocatorWhenRequestForLargeAllocationThenAllocatedSuccesfully) {
    constexpr size_t sharedIsaAllocationSize = MemoryConstants::pageSize2M * 2;
    constexpr size_t requestAllocationSize = sharedIsaAllocationSize * 2;

    auto &isaAllocator = pDevice->getIsaPoolAllocator();
    auto allocation = isaAllocator.requestGraphicsAllocationForIsa(false, requestAllocationSize);
    verifySharedIsaAllocation(allocation, 0, requestAllocationSize);
    isaAllocator.freeSharedIsaAllocation(allocation);
}

TEST_F(IsaPoolAllocatorTest, Given2MBLocalMemAlignmentEnabledWhenAllocatingIsaThenISAPoolSizeIsAlignedTo2MB) {
    auto mockProductHelper = new MockProductHelper;
    pDevice->getRootDeviceEnvironmentRef().productHelper.reset(mockProductHelper);
    mockProductHelper->is2MBLocalMemAlignmentEnabledResult = true;

    MockISAPoolAllocator mockIsaAllocator(pDevice);

    constexpr size_t requestAllocationSize = MemoryConstants::pageSize + 1;
    auto allocation = mockIsaAllocator.requestGraphicsAllocationForIsa(true, requestAllocationSize);

    EXPECT_NE(nullptr, allocation);
    ASSERT_GE(mockIsaAllocator.bufferPools.size(), 1u);

    auto *bufferPool = static_cast<MockISAPool *>(&mockIsaAllocator.bufferPools[0]);
    auto chunkAllocatorSize = bufferPool->chunkAllocator->getLeftSize() + bufferPool->chunkAllocator->getUsedSize();

    EXPECT_TRUE(isAligned(chunkAllocatorSize, MemoryConstants::pageSize2M));

    mockIsaAllocator.freeSharedIsaAllocation(allocation);
}