/*
 * Copyright (C) 2017-2020 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "shared/test/unit_test/mocks/mock_device.h"

#include "opencl/test/unit_test/fixtures/execution_model_kernel_fixture.h"
#include "opencl/test/unit_test/mocks/mock_kernel.h"
#include "opencl/test/unit_test/mocks/mock_program.h"
#include "test.h"

#include <memory>

using namespace NEO;

class MockKernelWithArgumentAccess : public Kernel {
  public:
    std::vector<SimpleKernelArgInfo> &getKernelArguments() {
        return kernelArguments;
    }

    class ObjectCountsPublic : public Kernel::ObjectCounts {
    };

    MockKernelWithArgumentAccess(Program *programArg, const KernelInfoContainer &kernelInfoArg) : Kernel(programArg, kernelInfoArg, false) {
    }

    void getParentObjectCountsPublic(MockKernelWithArgumentAccess::ObjectCountsPublic &objectCount) {
        getParentObjectCounts(objectCount);
    }
};

TEST(ParentKernelTest, WhenArgsAddedThenObjectCountsAreIncremented) {
    KernelInfo info;
    MockClDevice *device = new MockClDevice{new MockDevice};
    MockProgram program(toClDeviceVector(*device));
    SPatchExecutionEnvironment environment = {};
    environment.HasDeviceEnqueue = 1;

    info.patchInfo.executionEnvironment = &environment;

    MockKernelWithArgumentAccess kernel(&program, MockKernel::toKernelInfoContainer(info, device->getRootDeviceIndex()));

    std::vector<Kernel::SimpleKernelArgInfo> &args = kernel.getKernelArguments();

    Kernel::SimpleKernelArgInfo argInfo;
    argInfo.type = Kernel::kernelArgType::SAMPLER_OBJ;
    args.push_back(argInfo);
    argInfo.type = Kernel::kernelArgType::IMAGE_OBJ;
    args.push_back(argInfo);

    MockKernelWithArgumentAccess::ObjectCountsPublic objectCounts;

    kernel.getParentObjectCountsPublic(objectCounts);

    EXPECT_EQ(1u, objectCounts.imageCount);
    EXPECT_EQ(1u, objectCounts.samplerCount);

    delete device;
}

TEST(ParentKernelTest, WhenPatchingBlocksSimdSizeThenPatchIsAppliedCorrectly) {
    MockClDevice device{new MockDevice};
    auto rootDeviceIndex = device.getRootDeviceIndex();
    MockContext context(&device);
    std::unique_ptr<MockParentKernel> parentKernel(MockParentKernel::create(context, true));
    MockProgram *program = (MockProgram *)parentKernel->mockProgram;

    parentKernel->patchBlocksSimdSize(rootDeviceIndex);

    void *blockSimdSize = ptrOffset(parentKernel->getCrossThreadData(rootDeviceIndex), parentKernel->getKernelInfo(rootDeviceIndex).childrenKernelsIdOffset[0].second);
    uint32_t *simdSize = reinterpret_cast<uint32_t *>(blockSimdSize);

    EXPECT_EQ(program->blockKernelManager->getBlockKernelInfo(0)->getMaxSimdSize(), *simdSize);
}

TEST(ParentKernelTest, GivenParentKernelWhenCheckingForDeviceEnqueueThenTrueIsReturned) {
    MockClDevice device{new MockDevice};
    MockContext context(&device);
    std::unique_ptr<MockParentKernel> parentKernel(MockParentKernel::create(context));

    auto rootDeviceIndex = device.getRootDeviceIndex();
    EXPECT_TRUE(parentKernel->getKernelInfo(rootDeviceIndex).hasDeviceEnqueue());
}

TEST(ParentKernelTest, GivenNormalKernelWhenCheckingForDeviceEnqueueThenFalseIsReturned) {
    MockClDevice device{new MockDevice};
    MockKernelWithInternals kernel(device);

    EXPECT_FALSE(kernel.kernelInfo.hasDeviceEnqueue());
}

TEST(ParentKernelTest, WhenInitializingParentKernelThenBlocksSimdSizeIsPatched) {
    MockClDevice device{new MockDevice};
    auto rootDeviceIndex = device.getRootDeviceIndex();
    MockContext context(&device);
    std::unique_ptr<MockParentKernel> parentKernel(MockParentKernel::create(context, true));
    MockProgram *program = (MockProgram *)parentKernel->mockProgram;

    parentKernel->initialize();

    void *blockSimdSize = ptrOffset(parentKernel->getCrossThreadData(rootDeviceIndex), parentKernel->getKernelInfo(rootDeviceIndex).childrenKernelsIdOffset[0].second);
    uint32_t *simdSize = reinterpret_cast<uint32_t *>(blockSimdSize);

    EXPECT_EQ(program->blockKernelManager->getBlockKernelInfo(0)->getMaxSimdSize(), *simdSize);
}

TEST(ParentKernelTest, WhenInitializingParentKernelThenPrivateMemoryForBlocksIsAllocated) {
    MockClDevice device{new MockDevice};
    MockContext context(&device);
    std::unique_ptr<MockParentKernel> parentKernel(MockParentKernel::create(context, true));
    MockProgram *program = (MockProgram *)parentKernel->mockProgram;

    uint32_t crossThreadOffsetBlock = 0;

    auto infoBlock = new KernelInfo();
    SPatchAllocateStatelessDefaultDeviceQueueSurface *allocateDeviceQueueBlock = new SPatchAllocateStatelessDefaultDeviceQueueSurface;
    allocateDeviceQueueBlock->DataParamOffset = crossThreadOffsetBlock;
    allocateDeviceQueueBlock->DataParamSize = 8;
    allocateDeviceQueueBlock->SurfaceStateHeapOffset = 0;
    allocateDeviceQueueBlock->Size = 8;
    infoBlock->patchInfo.pAllocateStatelessDefaultDeviceQueueSurface = allocateDeviceQueueBlock;

    crossThreadOffsetBlock += 8;

    SPatchAllocateStatelessEventPoolSurface *eventPoolBlock = new SPatchAllocateStatelessEventPoolSurface;
    eventPoolBlock->DataParamOffset = crossThreadOffsetBlock;
    eventPoolBlock->DataParamSize = 8;
    eventPoolBlock->EventPoolSurfaceIndex = 0;
    eventPoolBlock->Size = 8;
    infoBlock->patchInfo.pAllocateStatelessEventPoolSurface = eventPoolBlock;

    crossThreadOffsetBlock += 8;

    auto privateSurfaceBlock = std::make_unique<SPatchAllocateStatelessPrivateSurface>();
    privateSurfaceBlock->DataParamOffset = crossThreadOffsetBlock;
    privateSurfaceBlock->DataParamSize = 8;
    privateSurfaceBlock->Size = 8;
    privateSurfaceBlock->SurfaceStateHeapOffset = 0;
    privateSurfaceBlock->Token = 0;
    privateSurfaceBlock->PerThreadPrivateMemorySize = 1000;
    infoBlock->patchInfo.pAllocateStatelessPrivateSurface = privateSurfaceBlock.get();

    crossThreadOffsetBlock += 8;

    SPatchThreadPayload *threadPayloadBlock = new SPatchThreadPayload;
    threadPayloadBlock->LocalIDXPresent = 0;
    threadPayloadBlock->LocalIDYPresent = 0;
    threadPayloadBlock->LocalIDZPresent = 0;
    threadPayloadBlock->HeaderPresent = 0;
    threadPayloadBlock->Size = 128;

    infoBlock->patchInfo.threadPayload = threadPayloadBlock;

    SPatchExecutionEnvironment *executionEnvironmentBlock = new SPatchExecutionEnvironment;
    *executionEnvironmentBlock = {};
    executionEnvironmentBlock->HasDeviceEnqueue = 1;
    infoBlock->patchInfo.executionEnvironment = executionEnvironmentBlock;

    SPatchDataParameterStream *streamBlock = new SPatchDataParameterStream;
    streamBlock->DataParameterStreamSize = 0;
    streamBlock->Size = 0;
    infoBlock->patchInfo.dataParameterStream = streamBlock;

    SPatchBindingTableState *bindingTable = new SPatchBindingTableState;
    bindingTable->Count = 0;
    bindingTable->Offset = 0;
    bindingTable->Size = 0;
    bindingTable->SurfaceStateOffset = 0;
    infoBlock->patchInfo.bindingTableState = bindingTable;

    SPatchInterfaceDescriptorData *idData = new SPatchInterfaceDescriptorData;
    idData->BindingTableOffset = 0;
    idData->KernelOffset = 0;
    idData->Offset = 0;
    idData->SamplerStateOffset = 0;
    idData->Size = 0;
    infoBlock->patchInfo.interfaceDescriptorData = idData;

    infoBlock->patchInfo.pAllocateStatelessGlobalMemorySurfaceWithInitialization = nullptr;
    infoBlock->patchInfo.pAllocateStatelessConstantMemorySurfaceWithInitialization = nullptr;

    infoBlock->heapInfo.pDsh = (void *)new uint64_t[64];
    infoBlock->crossThreadData = new char[crossThreadOffsetBlock];

    program->blockKernelManager->addBlockKernelInfo(infoBlock);

    parentKernel->initialize();

    EXPECT_NE(nullptr, program->getBlockKernelManager()->getPrivateSurface(program->getBlockKernelManager()->getCount() - 1));
}

struct ParentKernelFromBinaryTest : public ExecutionModelKernelFixture {

    void SetUp() override {
        ExecutionModelKernelFixture::SetUp("simple_block_kernel", "simple_block_kernel");
    }
};

TEST_F(ParentKernelFromBinaryTest, GivenParentKernelWhenGettingInstructionHeapSizeForExecutionModelThenSizeIsGreaterThanZero) {
    EXPECT_TRUE(pKernel->isParentKernel);
    EXPECT_LT(0u, pKernel->getInstructionHeapSizeForExecutionModel());
}