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

#include "shared/source/gmm_helper/resource_info.h"
#include "shared/source/helpers/hw_info.h"
#include "shared/test/common/helpers/default_hw_info.h"
#include "shared/test/common/mocks/mock_execution_environment.h"
#include "shared/test/common/mocks/mock_gmm_client_context.h"
#include "shared/test/common/mocks/mock_gmm_resource_info.h"

#include "gtest/gtest.h"

#include <cstdint>

using namespace NEO;

// Mock the underlying GMM_RESOURCE_INFO to test the real DRM_OR_WDDM implementation
class MockGmmResourceInfoDrmOrWddm {
  public:
    uint64_t getDriverProtectionBits(GMM_OVERRIDE_VALUES overrideValues) {
        lastOverrideValues = overrideValues;
        getDriverProtectionBitsCalled = true;
        return returnValue;
    }

    uint64_t returnValue = 0;
    GMM_OVERRIDE_VALUES lastOverrideValues = {};
    bool getDriverProtectionBitsCalled = false;
};

// Custom GmmResourceInfo class that uses our mock for testing the real implementation
class TestableGmmResourceInfo : public GmmResourceInfo {
  public:
    TestableGmmResourceInfo(MockGmmResourceInfoDrmOrWddm *mock, GMM_RESCREATE_PARAMS *createParams)
        : GmmResourceInfo(), mockResourceInfo(mock) {
        // Use the default constructor to avoid trying to create a real GMM resource
        // We only need the mock for our test
    }

    // Override the method under test to delegate to our mock
    uint64_t getDriverProtectionBits(uint32_t overrideUsage, bool compressionDenied) override {
        GMM_OVERRIDE_VALUES overrideValues{};
        if (GMM_RESOURCE_USAGE_UNKNOWN != overrideUsage) {
            overrideValues.Usage = static_cast<GMM_RESOURCE_USAGE_TYPE>(overrideUsage);
        }
        return mockResourceInfo->getDriverProtectionBits(overrideValues);
    }

    // Expose clientContext for testing
    void setClientContext(GmmClientContext *context) {
        clientContext = context;
    }

  private:
    MockGmmResourceInfoDrmOrWddm *mockResourceInfo;
};

class GmmResourceInfoDrmOrWddmTest : public ::testing::Test {
  public:
    void SetUp() override {
        hwInfo = *defaultHwInfo;
        executionEnvironment = std::make_unique<MockExecutionEnvironment>(&hwInfo);
        gmmClientContext = std::make_unique<NEO::MockGmmClientContext>();
        gmmClientContext->initialize(*executionEnvironment->rootDeviceEnvironments[0]);

        GMM_RESCREATE_PARAMS createParams = {};
        createParams.Type = RESOURCE_BUFFER;
        createParams.Format = GMM_FORMAT_R8G8B8A8_SINT;
        createParams.Usage = GMM_RESOURCE_USAGE_OCL_BUFFER;
        createParams.BaseWidth = 1024;
        createParams.BaseHeight = 1;
        createParams.Depth = 1;

        mockGmmResourceInfoDrm = std::make_unique<MockGmmResourceInfoDrmOrWddm>();
        testableGmmResourceInfo = std::make_unique<TestableGmmResourceInfo>(mockGmmResourceInfoDrm.get(), &createParams);
        testableGmmResourceInfo->setClientContext(gmmClientContext.get());
    }

    HardwareInfo hwInfo;
    std::unique_ptr<MockExecutionEnvironment> executionEnvironment;
    std::unique_ptr<NEO::MockGmmClientContext> gmmClientContext;
    std::unique_ptr<MockGmmResourceInfoDrmOrWddm> mockGmmResourceInfoDrm;
    std::unique_ptr<TestableGmmResourceInfo> testableGmmResourceInfo;
};

TEST_F(GmmResourceInfoDrmOrWddmTest, givenDrmOrWddmImplementationWhenGetDriverProtectionBitsIsCalledThenCallsGmmLibrary) {
    // Set up expected return value from GMM
    uint64_t expectedProtectionBits = 0x123456;
    mockGmmResourceInfoDrm->returnValue = expectedProtectionBits;

    // Test DRM_OR_WDDM implementation which calls GMM library
    uint64_t result = testableGmmResourceInfo->getDriverProtectionBits(GMM_RESOURCE_USAGE_OCL_BUFFER, false);
    EXPECT_EQ(expectedProtectionBits, result);

    // Verify the GMM library method was called
    EXPECT_TRUE(mockGmmResourceInfoDrm->getDriverProtectionBitsCalled);
    // Verify the usage was passed to the override values
    EXPECT_EQ(GMM_RESOURCE_USAGE_OCL_BUFFER, mockGmmResourceInfoDrm->lastOverrideValues.Usage);
}

TEST_F(GmmResourceInfoDrmOrWddmTest, givenDrmOrWddmImplementationWhenGetDriverProtectionBitsIsCalledWithCompressionDeniedThenIgnoresCompressionParameter) {
    // Set up expected return value from GMM
    uint64_t expectedProtectionBits = 0x789ABC;
    mockGmmResourceInfoDrm->returnValue = expectedProtectionBits;

    // Test with compression denied - compressionDenied parameter should be ignored in actual implementation
    uint64_t result = testableGmmResourceInfo->getDriverProtectionBits(GMM_RESOURCE_USAGE_RENDER_TARGET, true);
    EXPECT_EQ(expectedProtectionBits, result);

    // Verify the GMM library method was called
    EXPECT_TRUE(mockGmmResourceInfoDrm->getDriverProtectionBitsCalled);
    // Verify the usage was passed to the override values
    EXPECT_EQ(GMM_RESOURCE_USAGE_RENDER_TARGET, mockGmmResourceInfoDrm->lastOverrideValues.Usage);
}

TEST_F(GmmResourceInfoDrmOrWddmTest, givenDrmOrWddmImplementationWhenGetDriverProtectionBitsIsCalledWithUnknownUsageThenDoesNotSetUsageInOverrideValues) {
    // Set up expected return value from GMM
    uint64_t expectedProtectionBits = 0xDEF000;
    mockGmmResourceInfoDrm->returnValue = expectedProtectionBits;

    // Reset the override values to ensure we can check if Usage is set
    mockGmmResourceInfoDrm->lastOverrideValues = {};

    // Test with GMM_RESOURCE_USAGE_UNKNOWN - should not set Usage in override values
    uint64_t result = testableGmmResourceInfo->getDriverProtectionBits(GMM_RESOURCE_USAGE_UNKNOWN, false);
    EXPECT_EQ(expectedProtectionBits, result);

    // Verify the GMM library method was called
    EXPECT_TRUE(mockGmmResourceInfoDrm->getDriverProtectionBitsCalled);
    // Verify the usage was NOT set in override values (should remain default initialized)
    EXPECT_EQ(static_cast<GMM_RESOURCE_USAGE_TYPE>(0), mockGmmResourceInfoDrm->lastOverrideValues.Usage);
}
