#include "kernel_config.cuh"
#include <cuda_runtime.h>
#include <gtest/gtest.h>
#include <numeric>
#include <vector>

// Kernel to test the get_thread_id() function
__global__ void test_get_thread_id_kernel(size_t *output, size_t n_elements) {
  size_t i = get_thread_id();
  if (i < n_elements) {
    output[i] = i;
  }
}

// Test fixture for kernel config tests
class KernelConfigTest : public ::testing::Test {
protected:
  void SetUp() override {
    // Set up any common resources for the tests
  }

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

TEST_F(KernelConfigTest, GetThreadId) {
  const size_t n_elements = 10000;
  KernelConfig config = get_launch_config(n_elements);

  size_t *d_output;
  cudaMalloc(&d_output, n_elements * sizeof(size_t));

  test_get_thread_id_kernel<<<config.blocks, config.threads>>>(d_output,
                                                               n_elements);

  std::vector<size_t> h_output(n_elements);
  cudaMemcpy(h_output.data(), d_output, n_elements * sizeof(size_t),
             cudaMemcpyDeviceToHost);

  cudaFree(d_output);

  std::vector<size_t> expected(n_elements);
  std::iota(expected.begin(), expected.end(), 0);

  ASSERT_EQ(h_output, expected);
}