// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

#include "lib/aqlprofile/pm4/spm_builder.h"
#include "lib/aqlprofile/pm4/cmd_builder.h"
#include "lib/aqlprofile/pm4/cmd_config.h"
#include "lib/aqlprofile/pm4/trace_config.h"
#include "lib/aqlprofile/def/gpu_block_info.h"

#include <gtest/gtest.h>
#include <gmock/gmock.h>

#include <cstring>
#include <memory>
#include <vector>

using namespace pm4_builder;
// using namespace aql_profile;

namespace spm_builder_tests
{
// Mock SpmBuilder class for testing
class MockSpmBuilder : public SpmBuilder
{
public:
    MOCK_METHOD(void,
                Begin,
                (CmdBuffer * cmd_buffer,
                 const SpmConfig*       config,
                 const counters_vector& counters_vec),
                (override));
    MOCK_METHOD(void, End, (CmdBuffer * cmd_buffer, const SpmConfig* config), (override));
};

class SpmBuilderTest : public ::testing::Test
{
protected:
    void SetUp() override
    {
        test_config_ = SpmConfig{};

        // Set up default SPM config
        test_config_.sampleRate       = 1000;
        test_config_.data_buffer_ptr  = test_buffer_.data();
        test_config_.data_buffer_size = test_buffer_.size() * sizeof(uint32_t);

        // Initialize agent info for creating concrete SpmBuilder
        agent_info_ = AgentInfo{};
        strncpy(agent_info_.name, "gfx90a", sizeof(agent_info_.name) - 1);
        strncpy(agent_info_.gfxip, "gfx90a", sizeof(agent_info_.gfxip) - 1);
        agent_info_.cu_num               = 104;
        agent_info_.se_num               = 8;
        agent_info_.xcc_num              = 1;
        agent_info_.shader_arrays_per_se = 2;
    }

    void TearDown() override
    {
        // Clean up any resources
    }

    SpmConfig             test_config_;
    std::vector<uint32_t> test_buffer_{1024, 0};  // 4KB buffer initialized with zeros
    AgentInfo             agent_info_;
    counters_vector       test_counters_;
};

// Test 1: Begin function with valid parameters
TEST_F(SpmBuilderTest, BeginWithValidParameters)
{
    // Create a mock SpmBuilder
    MockSpmBuilder mock_spm_builder;
    CmdBuffer      cmd_buffer;

    // Set up expectations - Begin should be called once with the provided parameters
    EXPECT_CALL(mock_spm_builder, Begin(&cmd_buffer, &test_config_, ::testing::Ref(test_counters_)))
        .Times(1);

    // Call Begin method
    mock_spm_builder.Begin(&cmd_buffer, &test_config_, test_counters_);

    // Verify that the command buffer is still valid after the call
    EXPECT_GE(cmd_buffer.DwSize(), 0);
}

// Test 2: End function with valid parameters
TEST_F(SpmBuilderTest, EndWithValidParameters)
{
    // Create a mock SpmBuilder
    MockSpmBuilder mock_spm_builder;
    CmdBuffer      cmd_buffer;

    // Set up expectations - End should be called once with the provided parameters
    EXPECT_CALL(mock_spm_builder, End(&cmd_buffer, &test_config_)).Times(1);

    // Call End method
    mock_spm_builder.End(&cmd_buffer, &test_config_);

    // Verify that the command buffer is still valid after the call
    EXPECT_GE(cmd_buffer.DwSize(), 0);
}

// Test 5: Begin and End sequence with mock
TEST_F(SpmBuilderTest, BeginEndSequenceWithMock)
{
    MockSpmBuilder mock_spm_builder;
    CmdBuffer      cmd_buffer;

    // Set up expectations for a complete Begin-End sequence
    ::testing::InSequence seq;
    EXPECT_CALL(mock_spm_builder, Begin(&cmd_buffer, &test_config_, ::testing::Ref(test_counters_)))
        .Times(1);
    EXPECT_CALL(mock_spm_builder, End(&cmd_buffer, &test_config_)).Times(1);

    // Execute the sequence
    mock_spm_builder.Begin(&cmd_buffer, &test_config_, test_counters_);
    mock_spm_builder.End(&cmd_buffer, &test_config_);

    // Verify buffer state after complete sequence
    EXPECT_GE(cmd_buffer.DwSize(), 0);
}

// Test 6: Null parameter handling (defensive programming)
TEST_F(SpmBuilderTest, NullParameterHandling)
{
    MockSpmBuilder mock_spm_builder;

    // These tests verify that the mock can handle null parameters
    // In a real implementation, these should be handled gracefully or throw exceptions

    // Test with null command buffer - should be handled by implementation
    EXPECT_CALL(mock_spm_builder, Begin(nullptr, &test_config_, ::testing::Ref(test_counters_)))
        .Times(1);
    mock_spm_builder.Begin(nullptr, &test_config_, test_counters_);

    // Test with null config - should be handled by implementation
    CmdBuffer cmd_buffer;
    EXPECT_CALL(mock_spm_builder, Begin(&cmd_buffer, nullptr, ::testing::Ref(test_counters_)))
        .Times(1);
    mock_spm_builder.Begin(&cmd_buffer, nullptr, test_counters_);
}

}  // namespace spm_builder_tests
