#include "gpu.hpp"
#include "numeric_types/half.hpp"
#include <array>
#include <cassert>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <future>
#include <vector>

using namespace gpu;
using namespace std::chrono;

// WGSL Kernels

// Kernel to unpack 4x int8 (packed in i32) to 4x int32
const char *kPackedInt8ToInt32Kernel = R"(
  @group(0) @binding(0) var<storage, read_write> packed_input: array<i32>;
  @group(0) @binding(1) var<storage, read_write> unpacked_output: array<i32>;
  
  // Function to sign-extend an 8-bit value (represented in the lower bits of an i32)
  fn sign_extend_i8(val: i32) -> i32 {
    return (val << 24) >> 24;
  }
  
  @compute @workgroup_size({{workgroupSize}})
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let packed_idx: u32 = gid.x;
  
    // Check bounds for the PACKED input array
    if (packed_idx >= arrayLength(&packed_input)) {
      return;
    }
  
    let packed_val = packed_input[packed_idx];
  
    // Unpack and write 4 separate i32 values
    // Ensure the output buffer is large enough (4x the packed size)
    let base_output_idx = packed_idx * 4u;
  
    // Check bounds for the UNPACKED output array (optional but safer)
    // This assumes arrayLength(&unpacked_output) is at least 4 * arrayLength(&packed_input)
    if ((base_output_idx + 3u) >= arrayLength(&unpacked_output)) {
        return; // Avoid out-of-bounds write if something is wrong
    }
  
    unpacked_output[base_output_idx + 0u] = sign_extend_i8((packed_val >> 0u) & 0xFF);
    unpacked_output[base_output_idx + 1u] = sign_extend_i8((packed_val >> 8u) & 0xFF);
    unpacked_output[base_output_idx + 2u] = sign_extend_i8((packed_val >> 16u) & 0xFF);
    unpacked_output[base_output_idx + 3u] = sign_extend_i8((packed_val >> 24u) & 0xFF);
  }
  )";

// Kernel to pack 4x int32 back into 1x int32 (taking lower 8 bits)
const char *kInt32ToPackedInt8Kernel = R"(
  @group(0) @binding(0) var<storage, read_write> unpacked_input: array<i32>;
  @group(0) @binding(1) var<storage, read_write> packed_output: array<i32>;
  
  @compute @workgroup_size({{workgroupSize}})
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let packed_idx: u32 = gid.x; // Index for the PACKED output array
  
    // Check bounds for the PACKED output array
     if (packed_idx >= arrayLength(&packed_output)) {
      return;
    }
  
    let base_input_idx = packed_idx * 4u;
  
    // Check bounds for the UNPACKED input array (optional but safer)
    // Assumes arrayLength(&unpacked_input) is at least 4 * arrayLength(&packed_output)
     if ((base_input_idx + 3u) >= arrayLength(&unpacked_input)) {
        // Handle potential error or incomplete data - maybe write 0?
        packed_output[packed_idx] = 0;
        return;
    }
  
    // Read 4 separate i32 values
    let val0 = unpacked_input[base_input_idx + 0u];
    let val1 = unpacked_input[base_input_idx + 1u];
    let val2 = unpacked_input[base_input_idx + 2u];
    let val3 = unpacked_input[base_input_idx + 3u];
  
    // Pack the lower 8 bits of each into one i32
    var packed_result: i32 = 0;
    packed_result = packed_result | ((val0 & 0xFF) << 0u);
    packed_result = packed_result | ((val1 & 0xFF) << 8u);
    packed_result = packed_result | ((val2 & 0xFF) << 16u);
    packed_result = packed_result | ((val3 & 0xFF) << 24u);
  
    packed_output[packed_idx] = packed_result;
  }
  )";

// Simple addition kernel for i32
const char *kSimpleAddKernelI32 = R"(
  @group(0) @binding(0) var<storage, read_write> a: array<{{precision}}>;
  @group(0) @binding(1) var<storage, read_write> b: array<{{precision}}>;
  @group(0) @binding(2) var<storage, read_write> c: array<{{precision}}>;
  
  @compute @workgroup_size({{workgroupSize}})
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i: u32 = gid.x;
    if (i < arrayLength(&a)) {
      c[i] = a[i] + b[i];
    }
  }
  )";

// A simple WGSL copy kernel that copies input to output.
static const char *kCopyKernel = R"(
  @group(0) @binding(0) var<storage, read_write> inp: array<{{precision}}>;
  @group(0) @binding(1) var<storage, read_write> out: array<{{precision}}>;
  @compute @workgroup_size({{workgroupSize}})
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i: u32 = gid.x;
    if (i < arrayLength(&inp)) {
      out[i] = inp[i];
    }
  }
  )";

// Forward declarations:
void testToCPUWithTensor();
void testToCPUWithBuffer();
void testToCPUWithTensorSourceOffset();
void testToCPUWithBufferSourceOffset();
void stressTestToCPU();
void testToCPUWithHalf();
void testToCPUWithFloat();
void testToCPUWithDouble();
void testToCPUWithint8();
void testToCPUWithint16();
void testToCPUWithint();
void testToCPUWithint64();
void testToCPUWithUint8();
void testToCPUWithUint16();
void testToCPUWithUint32();
void testToCPUWithUint64();
void testNumTypeSizes();
void testToCPUUnpack();
void testCopyShaderPackedUnpack_int8();
void testAddKernelInt8();

int main() {
  LOG(kDefLog, kInfo, "Running GPU integration tests...");
  testAddKernelInt8();
  testCopyShaderPackedUnpack_int8();
  testToCPUUnpack();
  testToCPUWithTensor();
  testToCPUWithBuffer();
  testToCPUWithTensorSourceOffset();
  testToCPUWithBufferSourceOffset();
  testToCPUWithHalf();
  testToCPUWithFloat();
  testToCPUWithDouble();
  testToCPUWithint8();
  testToCPUWithint16();
  testToCPUWithint();
  testToCPUWithint64();
  testToCPUWithUint8();
  testToCPUWithUint16();
  testToCPUWithUint32();
  testToCPUWithUint64();
  testNumTypeSizes();
  stressTestToCPU();
  testHalf();
  LOG(kDefLog, kInfo, "All tests passed.");
  return 0;
}

void testAddKernelInt8() {
  LOG(kDefLog, kInfo, "Running testAddKernelInt8 (with conversion kernels)...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024; // Logical number of int8 elements
  std::vector<int8_t> aInput(N), bInput(N), result(N);
  std::vector<int8_t> expected(N);

  // CPU Data Setup
  for (size_t i = 0; i < N; ++i) {
    // Values in range [-10, 9]
    aInput[i] = static_cast<int8_t>((i % 20) - 10);
    bInput[i] = static_cast<int8_t>(((2 * i) % 20) - 10);
    // Compute expected as int then cast back.
    int temp = static_cast<int>(aInput[i]) + static_cast<int>(bInput[i]);
    expected[i] = static_cast<int8_t>(temp);
    result[i] = 0;
  }

  // These store the int8 data packed into i32 format on the GPU
  Tensor aTensorPacked = createTensor(ctx, Shape{N}, ki8, (const int8_t *)aInput.data());
  Tensor bTensorPacked = createTensor(ctx, Shape{N}, ki8, (const int8_t *)bInput.data());
  // Final output tensor, also in packed format
  Tensor outputTensorPacked = createTensor(ctx, Shape{N}, ki8);

  // These will hold the data converted to one i32 per original int8 element
  Tensor aTensorUnpacked = createTensor(ctx, Shape{N}, ki32);
  Tensor bTensorUnpacked = createTensor(ctx, Shape{N}, ki32);
  Tensor outputTensorUnpacked =
      createTensor(ctx, Shape{N}, ki32); // For the simple add result

  constexpr uint32_t workgroupSize = 256;
  size_t packedCount = (N + 3) / 4; // Number of i32 elements in packed buffers
  size_t unpackedCount = N; // Number of i32 elements in unpacked buffers

  // Convert Packed Inputs to Unpacked i32
  Kernel unpackKernelA =
      createKernel(ctx, {kPackedInt8ToInt32Kernel, workgroupSize, ki32},
                   Bindings{aTensorPacked, aTensorUnpacked},
                   {cdiv(packedCount, workgroupSize), 1,
                    1}); // Dispatch based on packed size
  Kernel unpackKernelB =
      createKernel(ctx, {kPackedInt8ToInt32Kernel, workgroupSize, ki32},
                   Bindings{bTensorPacked, bTensorUnpacked},
                   {cdiv(packedCount, workgroupSize), 1, 1});
  // Dispatch based on packed size
  dispatchKernel(ctx, unpackKernelA);
  dispatchKernel(ctx, unpackKernelB);

  // Perform Simple Addition on Unpacked i32
  Kernel simpleAddKernel = createKernel(
      ctx, {kSimpleAddKernelI32, workgroupSize, ki32},
      Bindings{aTensorUnpacked, bTensorUnpacked, outputTensorUnpacked},
      {cdiv(unpackedCount, workgroupSize), 1,
       1}); // Dispatch based on unpacked size
  dispatchKernel(ctx, simpleAddKernel);

  // Convert Unpacked i32 Result back to Packed
  Kernel packKernel =
      createKernel(ctx, {kInt32ToPackedInt8Kernel, workgroupSize, ki32},
                   Bindings{outputTensorUnpacked, outputTensorPacked},
                   {cdiv(packedCount, workgroupSize), 1,
                    1}); // Dispatch based on packed size
  dispatchKernel(ctx, packKernel);

  // Copy Final Packed Result to CPU and Unpack
  // Use the original toCPU for ki8, which handles the final CPU-side unpacking
  toCPU(ctx, outputTensorPacked, ki8, result.data(), 0);

  for (size_t i = 0; i < N; ++i) {
    assert(result[i] == expected[i]);
  }

  LOG(kDefLog, kInfo, "testAddKernelInt8 (with conversion kernels) passed.");
}

void testCopyShaderPackedUnpack_int8() {
  LOG(kDefLog, kInfo, "Running testCopyShaderPackedUnpack_int8...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::vector<int8_t> inputData(N), outputData(N);
  for (size_t i = 0; i < N; ++i) {
    // Values between -128 and 127.
    inputData[i] = static_cast<int8_t>((i % 256) - 128);
  }

  // Create an input tensor using the int8_t overload.
  // Under the hood the data is packed into int32_t.
  Tensor inputTensor = createTensor(ctx, Shape{N}, ki8, inputData.data());

  // Create an output tensor of the same shape and unsupported type.
  Tensor outputTensor = createTensor(ctx, Shape{N}, ki8);

  // Our copy shader (kCopyKernel) expects to work with supported types.
  // Since int8_t is packed into int32_t, we pass 'ki32' as our shader
  // precision.
  Kernel copyKernel =
      createKernel(ctx, {kCopyKernel, 256, ki32},
                   Bindings{inputTensor, outputTensor}, {cdiv(N, 256), 1, 1});
  dispatchKernel(ctx, copyKernel);

  // Now retrieve the output from the GPU and unpack from the packed int32_t
  // back to int8_t.
  toCPU(ctx, outputTensor, ki8, outputData.data(), 0);

  // Verify the unpacked data matches the original input.
  for (size_t i = 0; i < N; ++i) {
    assert(inputData[i] == outputData[i]);
  }
  LOG(kDefLog, kInfo, "testCopyShaderPackedUnpack_int8 passed.");
}

void testToCPUUnpack() {
  LOG(kDefLog, kInfo, "Running testToCPUUnpack...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  // Test for double (kf64 -> packed as kf32)
  {
    constexpr size_t N = 1024;
    std::vector<double> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<double>(i) * 3.14;
    }
    Tensor tensor = createTensor(ctx, Shape{N}, kf64, inputData.data());
    toCPU(ctx, tensor, kf64, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      // Allow for a very small epsilon error due to float conversion.
      assert(fabs(inputData[i] - outputData[i]) < 1e-4);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for double passed.");
  }

  // Test for int8_t (ki8 -> packed as ki32)
  {
    constexpr size_t N = 1024;
    std::vector<int8_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<int8_t>((i % 256) - 128);
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ki8, inputData.data());
    toCPU(ctx, tensor, ki8, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for int8_t passed.");
  }

  // Test for int16_t (ki16 -> packed as ki32)
  {
    constexpr size_t N = 1024;
    std::vector<int16_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<int16_t>((i % 65536) - 32768);
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ki16, inputData.data());
    toCPU(ctx, tensor, ki16, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for int16_t passed.");
  }

  // Test for int64_t (ki64 -> packed as two ki32s)
  {
    constexpr size_t N = 1024;
    std::vector<int64_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<int64_t>(i) - 512;
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ki64, inputData.data());
    toCPU(ctx, tensor, ki64, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for int64_t passed.");
  }

  // Test for uint8_t (ku8 -> packed as ku32)
  {
    constexpr size_t N = 1024;
    std::vector<uint8_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<uint8_t>(i % 256);
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ku8, inputData.data());
    toCPU(ctx, tensor, ku8, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for uint8_t passed.");
  }

  // Test for uint16_t (ku16 -> packed as ku32)
  {
    constexpr size_t N = 1024;
    std::vector<uint16_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<uint16_t>(i % 65536);
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ku16, inputData.data());
    toCPU(ctx, tensor, ku16, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for uint16_t passed.");
  }

  // Test for uint64_t (ku64 -> packed as two ku32s)
  {
    constexpr size_t N = 1024;
    std::vector<uint64_t> inputData(N), outputData(N);
    for (size_t i = 0; i < N; ++i) {
      inputData[i] = static_cast<uint64_t>(i) * 123456789ULL;
    }
    Tensor tensor = createTensor(ctx, Shape{N}, ku64, inputData.data());
    toCPU(ctx, tensor, ku64, outputData.data(), 0);
    for (size_t i = 0; i < N; ++i) {
      assert(inputData[i] == outputData[i]);
    }
    LOG(kDefLog, kInfo, "toCPUUnpack for uint64_t passed.");
  }

  LOG(kDefLog, kInfo, "All toCPUUnpack tests passed.");
}

void testNumTypeSizes() {
  LOG(kDefLog, kInfo, "Running testNumTypeSizes...");

  assert(sizeBytes(kf16) == 2);
  assert(sizeBytes(kf32) == 4);
  assert(sizeBytes(ki8) == sizeof(uint32_t));   // ki8 is packed into uint32_t.
  assert(sizeBytes(ki16) == sizeof(uint32_t)); // ki16 is packed into uint32_t.
  assert(sizeBytes(ki32) == sizeof(int32_t));  // typically 4
  assert(sizeBytes(ku8) == sizeof(uint32_t));  // ku8 is packed into uint32_t.
  assert(sizeBytes(ku16) == sizeof(uint32_t)); // ku16 is packed into uint32_t.
  assert(sizeBytes(ku32) == sizeof(uint32_t)); // typically 4

  LOG(kDefLog, kInfo, "testNumTypeSizes passed.");
}

// Test using half-precision (16-bit float) data.
void testToCPUWithHalf() {
  LOG(kDefLog, kInfo, "Running testToCPUWithHalf...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<half, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    // Construct half from float.
    inputData[i] = half(static_cast<float>(i));
  }

  Tensor inputTensor = createTensor(ctx, Shape{N}, kf16, inputData.data());

  // Copy GPU output to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy (using float conversion for approximate equality).
  for (size_t i = 0; i < N; ++i) {
    float inVal = static_cast<float>(inputData[i]);
    float outVal = static_cast<float>(outputData[i]);
    // Use a small epsilon to compare half values.
    assert(fabs(inVal - outVal) <= 0.01f);
  }
  LOG(kDefLog, kInfo, "testToCPUWithHalf passed.");
}

// Test using float (32-bit) data.
void testToCPUWithFloat() {
  LOG(kDefLog, kInfo, "Running testToCPUWithFloat...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<float, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<float>(i * 1.5f);
    outputData[i] = 0.0f;
  }

  Tensor inputTensor = createTensor(ctx, Shape{N}, kf32, inputData.data());

  // Copy GPU output to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    assert(inputData[i] == outputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithFloat passed.");
}

// Test using double (64-bit floating point) data.
void testToCPUWithDouble() {
  LOG(kDefLog, kInfo, "Running testToCPUWithDouble...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<double, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<double>(i) * 2.5;
    outputData[i] = 0.0;
  }

  Tensor inputTensor = createTensor(ctx, Shape{N}, kf64, inputData.data());

  // Copy GPU output to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    assert(inputData[i] == outputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithDouble passed.");
}

void testToCPUWithint8() {
  LOG(kDefLog, kInfo, "Running testToCPUWithint8...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<int8_t, N> inputData, outputData;
  // Use a range that includes negative values.
  for (size_t i = 0; i < N; ++i) {
    // Values between -128 and 127.
    inputData[i] = static_cast<int8_t>((i % 256) - 128);
    outputData[i] = 0;
  }

  // Create a tensor for int8_t.
  Tensor inputTensor = createTensor(ctx, Shape{N}, ki8, inputData.data());

  // Synchronously copy the GPU tensor data to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %d", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %d", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithint8 passed.");
}

// Test using int16_t data.
void testToCPUWithint16() {
  LOG(kDefLog, kInfo, "Running testToCPUWithint16...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<int16_t, N> inputData, outputData;
  // Use a range that includes negative values.
  for (size_t i = 0; i < N; ++i) {
    // Values between -32768 and 32767.
    inputData[i] = static_cast<int16_t>((i % 65536) - 32768);
    outputData[i] = 0;
  }

  // Create a tensor for int16_t.
  Tensor inputTensor = createTensor(ctx, Shape{N}, ki16, inputData.data());

  // Synchronously copy the GPU tensor data to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %d", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %d", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithint16 passed.");
}

// Test using int (int32_t) data.
void testToCPUWithint() {
  LOG(kDefLog, kInfo, "Running testToCPUWithint...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<int32_t, N> inputData, outputData;
  // Fill with sample data.
  for (size_t i = 0; i < N; ++i) {
    inputData[i] =
        static_cast<int32_t>(i - 512); // Negative and positive values.
    outputData[i] = 0;
  }

  // Create a tensor for int32_t.
  Tensor inputTensor = createTensor(ctx, Shape{N}, ki32, inputData.data());

  // Synchronously copy the GPU tensor data to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %d", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %d", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithint passed.");
}

// Test using int64_t (64-bit signed integer) data.
void testToCPUWithint64() {
  LOG(kDefLog, kInfo, "Running testToCPUWithint64...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<int64_t, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] =
        static_cast<int64_t>(i) - 512; // Some negative and positive values.
    outputData[i] = 0;
  }

  Tensor inputTensor = createTensor(ctx, Shape{N}, ki64, inputData.data());

  // Copy GPU output to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    assert(inputData[i] == outputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithint64 passed.");
}

void testToCPUWithUint8() {
  LOG(kDefLog, kInfo, "Running testToCPUWithUint8...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<uint8_t, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<uint8_t>(i % 256);
    outputData[i] = 0;
  }

  Tensor inputTensor = createTensor(
      ctx, Shape{N}, ku8, reinterpret_cast<const uint8_t *>(inputData.data()));

  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Verify the output matches the input.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %u", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %u", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithUint8 passed.");
}

void testToCPUWithUint16() {
  LOG(kDefLog, kInfo, "Running testToCPUWithUint16...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<uint16_t, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<uint16_t>(i % 65536);
    outputData[i] = 0;
  }

  Tensor inputTensor =
      createTensor(ctx, Shape{N}, ku16,
                   reinterpret_cast<const uint16_t *>(inputData.data()));

  // Synchronously copy GPU output to CPU using the tensor overload.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Verify the output matches the input.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %u", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %u", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithUint16 passed.");
}

void testToCPUWithUint32() {
  LOG(kDefLog, kInfo, "Running testToCPUWithUint32...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<uint32_t, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<uint32_t>(i);
    outputData[i] = 0;
  }

  Tensor inputTensor =
      createTensor(ctx, Shape{N}, ku32,
                   reinterpret_cast<const uint32_t *>(inputData.data()));

  // Synchronously copy GPU output to CPU using the tensor overload.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Verify the output matches the input.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %u", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %u", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithUint32 passed.");
}

// Test using uint64_t (64-bit unsigned integer) data.
void testToCPUWithUint64() {
  LOG(kDefLog, kInfo, "Running testToCPUWithUint64...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<uint64_t, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<uint64_t>(i);
    outputData[i] = 0;
  }

  // Assuming a new NumType 'ku64' for 64-bit unsigned integers.
  Tensor inputTensor = createTensor(ctx, Shape{N}, ku64, inputData.data());

  // Copy GPU output to CPU.
  toCPU(ctx, inputTensor, outputData.data(), sizeof(outputData));

  // Validate the copy.
  for (size_t i = 0; i < N; ++i) {
    assert(inputData[i] == outputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithUint64 passed.");
}

// Test using the overload that takes a Tensor.
void testToCPUWithTensor() {
  LOG(kDefLog, kInfo, "Running testToCPUWithTensor...");

// Create a real GPU context.
#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<float, N> inputData, outputData;
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<float>(i);
    outputData[i] = 0.0f;
  }

  // Create input and output tensors.
  Tensor inputTensor = createTensor(ctx, Shape{N}, kf32, inputData.data());
  Tensor outputTensor = createTensor(ctx, Shape{N}, kf32);

  // Create and dispatch the copy kernel.
  Kernel copyKernel =
      createKernel(ctx, {kCopyKernel, 256, kf32},
                   Bindings{inputTensor, outputTensor}, {cdiv(N, 256), 1, 1});
  dispatchKernel(ctx, copyKernel);

  // Synchronously copy GPU output to CPU using the tensor overload.
  toCPU(ctx, outputTensor, outputData.data(), sizeof(outputData));

  // Verify the output matches the input.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "inputData[%zu] = %f", i, inputData[i]);
    // LOG(kDefLog, kInfo, "outputData[%zu] = %f", i, outputData[i]);
    assert(outputData[i] == inputData[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithTensor passed.");
}

// Test using the overload that takes a raw GPU buffer.
// We reuse the Tensor's underlying buffer for this test.
void testToCPUWithBuffer() {
  LOG(kDefLog, kInfo, "Running testToCPUWithBuffer...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  std::array<float, N> data, outputData;
  for (size_t i = 0; i < N; ++i) {
    data[i] = static_cast<float>(i * 2);
    outputData[i] = 0.0f;
  }

  // Create a tensor to allocate a GPU buffer and initialize it.
  Tensor tensor = createTensor(ctx, Shape{N}, kf32, data.data());

  // Now extract the raw GPU buffer from the tensor.
  WGPUBuffer gpuBuffer = tensor.data.buffer;

  // Use the WGPUBuffer overload. This call returns a future.
  auto future =
      toCPUAsync(ctx, gpuBuffer, outputData.data(), sizeof(outputData), 0);
  wait(ctx, future);

  // Verify that the CPU output matches the original data.
  for (size_t i = 0; i < N; ++i) {
    // LOG(kDefLog, kInfo, "outputData[%zu] = %f", i, outputData[i]);
    assert(outputData[i] == data[i]);
  }
  LOG(kDefLog, kInfo, "testToCPUWithBuffer passed.");
}

void testToCPUWithTensorSourceOffset() {
  LOG(kDefLog, kInfo, "Running testToCPUWithTensorSourceOffset...");
#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t numElements = 25;
  constexpr size_t sourceOffsetElements = 5; // Skip first 5 elements
  constexpr size_t copyCount = 10;           // Number of floats to copy
  size_t copySize = copyCount * sizeof(float);

  // Create an input array with known data.
  std::array<float, numElements> inputData{};
  for (size_t i = 0; i < numElements; ++i) {
    inputData[i] = static_cast<float>(i + 50); // Arbitrary values
  }
  // Create a tensor from the full data.
  Tensor tensor = createTensor(ctx, Shape{numElements}, kf32, inputData.data());

  // Allocate a destination CPU buffer exactly as large as the data we want to
  // copy.
  std::vector<float> cpuOutput(copyCount, -1.0f);

  // Set sourceOffset to skip the first few float elements
  size_t sourceOffsetBytes = sourceOffsetElements * sizeof(float);
  // Call the tensor overload with sourceOffset and destOffset = 0.
  auto future =
      toCPUAsync(ctx, tensor, cpuOutput.data(), copySize, sourceOffsetBytes);
  wait(ctx, future);

  // Verify the copied data matches the expected subset.
  for (size_t i = 0; i < copyCount; ++i) {
    float expected = inputData[sourceOffsetElements + i];
    float actual = cpuOutput[i];
    // LOG(kDefLog, kInfo, "cpuOutput[%zu] = %f", i, actual);
    // LOG(kDefLog, kInfo, "expected[%zu] = %f", i, expected);
    assert(expected == actual);
  }
  LOG(kDefLog, kInfo, "testToCPUWithTensorSourceOffset passed.");
}

void testToCPUWithBufferSourceOffset() {
  LOG(kDefLog, kInfo, "Running testToCPUWithBufferSourceOffset...");
#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t numElements = 30;
  constexpr size_t sourceOffsetElements = 7; // Skip first 7 elements
  constexpr size_t copyCount = 12;           // Number of floats to copy
  size_t copySize = copyCount * sizeof(float);

  // Create an input array with arbitrary data.
  std::array<float, numElements> inputData{};
  for (size_t i = 0; i < numElements; ++i) {
    inputData[i] = static_cast<float>(i + 100);
  }
  // Create a tensor to initialize a GPU buffer.
  Tensor tensor = createTensor(ctx, Shape{numElements}, kf32, inputData.data());
  // Extract the raw GPU buffer from the tensor.
  WGPUBuffer buffer = tensor.data.buffer;

  // Allocate a destination CPU buffer exactly as large as needed.
  std::vector<float> cpuOutput(copyCount, -2.0f);
  size_t sourceOffsetBytes = sourceOffsetElements * sizeof(float);

  // Call the buffer overload with sourceOffset and destOffset = 0.
  auto future =
      toCPUAsync(ctx, buffer, cpuOutput.data(), copySize, sourceOffsetBytes);
  wait(ctx, future);

  // Verify that the copied data matches the expected subset.
  for (size_t i = 0; i < copyCount; ++i) {
    float expected = inputData[sourceOffsetElements + i];
    float actual = cpuOutput[i];
    // LOG(kDefLog, kInfo, "cpuOutput[%zu] = %f", i, actual);
    // LOG(kDefLog, kInfo, "expected[%zu] = %f", i, expected);
    assert(expected == actual);
  }
  LOG(kDefLog, kInfo, "testToCPUWithBufferSourceOffset passed.");
}

void stressTestToCPU() {
  LOG(kDefLog, kInfo, "Running stressTestToCPU for 2 seconds...");

#ifdef USE_DAWN_API
  Context ctx = createContextByGpuIdx(0);
#else
  Context ctx = createContext();
#endif

  constexpr size_t N = 1024;
  // Create a persistent tensor with some test data.
  std::vector<float> inputData(N, 0.0f);
  for (size_t i = 0; i < N; ++i) {
    inputData[i] = static_cast<float>(i);
  }
  Tensor tensor = createTensor(ctx, Shape{N}, kf32, inputData.data());

  // Prepare to run for one second.
  auto startTime = high_resolution_clock::now();
  size_t opCount = 0;
  while (high_resolution_clock::now() - startTime < seconds(2)) {
    // Allocate an output buffer (using a shared_ptr so it stays valid until the
    // future completes)
    auto outputData = std::make_shared<std::vector<float>>(N, 0.0f);
    // Use the tensor overload; we’re copying the entire tensor (destOffset = 0)
    // log count
    auto fut =
        toCPUAsync(ctx, tensor, outputData->data(), N * sizeof(float), 0);
    wait(ctx, fut);
    ++opCount;
  }

  auto endTime = high_resolution_clock::now();
  auto totalMs = duration_cast<milliseconds>(endTime - startTime).count();
  double throughput = (opCount / (totalMs / 1000.0));

  LOG(kDefLog, kInfo,
      "Stress test completed:\n"
      "  %zu GPU to CPU operations in %lld ms\n"
      "  Throughput: %.2f ops/sec",
      opCount, totalMs, throughput);
}