#include <gtest/gtest.h>
#include "../trace_config.h"

namespace pm4_builder {

class TraceConfigTest : public ::testing::Test {
protected:
  void SetUp() override {
    // Setup default configuration
    config.sampleRate = 1000;
    config.se_number = 4;  // Use se_number instead of spm_se_number_total
    config.se_mask = 0x0F;  // All 4 SEs enabled
    config.capacity_per_se = 0x2000;
    config.capacity_per_disabled_se = 0x1000;
  }

  TraceConfig config;
};

TEST_F(TraceConfigTest, DefaultValues) {
  TraceConfig default_config;
  
  // Check default initialization values
  EXPECT_EQ(default_config.targetCu, 0);
  EXPECT_EQ(default_config.vmIdMask, 0);
  EXPECT_EQ(default_config.simd_sel, 0xF);
  EXPECT_EQ(default_config.sampleRate, 625);
  EXPECT_EQ(default_config.perfMASK, ~0u);
  EXPECT_TRUE(default_config.spm_sq_32bit_mode);
  EXPECT_FALSE(default_config.spm_has_core1);
  EXPECT_EQ(default_config.se_mask, 0x11);
}

TEST_F(TraceConfigTest, SEConfiguration) {
  // Configure SE target CUs and base addresses
  config.target_cu_per_se[0] = 2;   // SE0: CU2
  config.target_cu_per_se[1] = -1;  // SE1: disabled
  config.target_cu_per_se[2] = 4;   // SE2: CU4
  config.target_cu_per_se[3] = 1;   // SE3: CU1

  config.se_base_addresses[0] = 0x1000;
  config.se_base_addresses[1] = 0x2000;
  config.se_base_addresses[2] = 0x3000;
  config.se_base_addresses[3] = 0x4000;

  // Test target CU retrieval
  EXPECT_EQ(config.GetTargetCU(0), 2);
  EXPECT_EQ(config.GetTargetCU(1), -1);
  EXPECT_EQ(config.GetTargetCU(2), 4);
  EXPECT_EQ(config.GetTargetCU(3), 1);

  // Test SE base address retrieval
  EXPECT_EQ(config.GetSEBaseAddr(0), 0x1000);
  EXPECT_EQ(config.GetSEBaseAddr(1), 0x2000);
  EXPECT_EQ(config.GetSEBaseAddr(2), 0x3000);
  EXPECT_EQ(config.GetSEBaseAddr(3), 0x4000);

  // Test SE capacity calculations
  EXPECT_EQ(config.GetCapacity(0), config.capacity_per_se);        // Enabled SE
  EXPECT_EQ(config.GetCapacity(1), config.capacity_per_disabled_se); // Disabled SE
  EXPECT_EQ(config.GetCapacity(2), config.capacity_per_se);        // Enabled SE
  EXPECT_EQ(config.GetCapacity(3), config.capacity_per_se);        // Enabled SE
}

TEST_F(TraceConfigTest, SEMaskConfiguration) {
  // Test different SE mask configurations
  config.se_mask = 0x5;  // Enable SE0 and SE2, disable SE1 and SE3

  // Setup target CUs
  config.target_cu_per_se[0] = 0;   // SE0 enabled
  config.target_cu_per_se[1] = -1;  // SE1 disabled
  config.target_cu_per_se[2] = 1;   // SE2 enabled
  config.target_cu_per_se[3] = -1;  // SE3 disabled

  EXPECT_EQ(config.GetSEmask(), 0x5);
  EXPECT_EQ(config.GetTargetCU(0), 0);
  EXPECT_EQ(config.GetTargetCU(1), -1);
  EXPECT_EQ(config.GetTargetCU(2), 1);
  EXPECT_EQ(config.GetTargetCU(3), -1);
}

TEST_F(TraceConfigTest, BufferConfiguration) {
  const size_t BUFFER_SIZE = 4096;
  char data_buffer[BUFFER_SIZE];
  char control_buffer[BUFFER_SIZE];

  // Configure buffers
  config.data_buffer_ptr = data_buffer;
  config.data_buffer_size = BUFFER_SIZE;
  config.control_buffer_ptr = control_buffer;
  config.control_buffer_size = BUFFER_SIZE;

  EXPECT_EQ(config.data_buffer_ptr, data_buffer);
  EXPECT_EQ(config.data_buffer_size, BUFFER_SIZE);
  EXPECT_EQ(config.control_buffer_ptr, control_buffer);
  EXPECT_EQ(config.control_buffer_size, BUFFER_SIZE);
}

TEST_F(TraceConfigTest, PerformanceConfiguration) {
  // Test performance counter configuration
  config.perfMASK = 0xF0F0;
  config.perfCTRL = 0x1234;
  
  // Add some performance counters
  config.perfcounters.push_back({0, 1});  // Counter 0, Instance 1
  config.perfcounters.push_back({2, 3});  // Counter 2, Instance 3

  EXPECT_EQ(config.perfMASK, 0xF0F0);
  EXPECT_EQ(config.perfCTRL, 0x1234);
  ASSERT_EQ(config.perfcounters.size(), 2);
  EXPECT_EQ(config.perfcounters[0].first, 0);
  EXPECT_EQ(config.perfcounters[0].second, 1);
  EXPECT_EQ(config.perfcounters[1].first, 2);
  EXPECT_EQ(config.perfcounters[1].second, 3);
}

TEST_F(TraceConfigTest, ConcurrentConfiguration) {
  // Test concurrent kernel configuration
  config.concurrent = 2;

  // Configure per-SE capacities for concurrent mode
  config.capacity_per_se = 0x4000;
  config.capacity_per_disabled_se = 0x2000;

  // Setup multiple SEs with different target CUs
  for (uint32_t se = 0; se < config.se_number; se++) {
    config.target_cu_per_se[se] = se % 2 ? -1 : se;  // Alternate between enabled/disabled
    config.se_base_addresses[se] = 0x1000 * (se + 1);
  }

  EXPECT_EQ(config.concurrent, 2);

  // Verify SE configuration in concurrent mode
  for (uint32_t se = 0; se < config.se_number; se++) {
    if (se % 2 == 0) {
      EXPECT_EQ(config.GetTargetCU(se), se);
      EXPECT_EQ(config.GetCapacity(se), config.capacity_per_se);
    } else {
      EXPECT_EQ(config.GetTargetCU(se), -1);
      EXPECT_EQ(config.GetCapacity(se), config.capacity_per_disabled_se);
    }
    EXPECT_EQ(config.GetSEBaseAddr(se), 0x1000 * (se + 1));
  }
}

TEST_F(TraceConfigTest, ExceptionHandling) {
  // Test accessing non-existent SE configurations
  EXPECT_THROW(config.GetTargetCU(99), std::out_of_range);
  EXPECT_THROW(config.GetSEBaseAddr(99), std::out_of_range);
}

} // namespace pm4_builder
