// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include <cstdint>
#include <vector>

#include "iree/hal/cts/util/test_base.h"

namespace iree::hal::cts {

using ::testing::ContainerEq;

class CommandBufferUpdateBufferTest : public CtsTestBase<> {
 protected:
  iree_hal_buffer_ref_t MakeTargetRef(iree_hal_buffer_t* buffer,
                                      iree_device_size_t offset,
                                      iree_device_size_t length) {
    if (recording_mode() == RecordingMode::kIndirect) {
      return iree_hal_make_indirect_buffer_ref(/*binding=*/0, offset, length);
    }
    return iree_hal_make_buffer_ref(buffer, offset, length);
  }

  iree_host_size_t binding_capacity() {
    return recording_mode() == RecordingMode::kIndirect ? 1 : 0;
  }

  iree_status_t SubmitWithBindings(iree_hal_command_buffer_t* command_buffer,
                                   iree_hal_buffer_t* target_buffer) {
    if (recording_mode() == RecordingMode::kIndirect) {
      const iree_hal_buffer_binding_t bindings[] = {
          {target_buffer, 0, IREE_HAL_WHOLE_BUFFER},
      };
      return SubmitCommandBufferAndWait(
          command_buffer,
          iree_hal_buffer_binding_table_t{IREE_ARRAYSIZE(bindings), bindings});
    }
    return SubmitCommandBufferAndWait(command_buffer);
  }
};

TEST_P(CommandBufferUpdateBufferTest, WholeBuffer) {
  iree_device_size_t target_buffer_size = 16;
  std::vector<uint8_t> source_buffer{0x01, 0x02, 0x03, 0x04,  //
                                     0x05, 0x06, 0x07, 0x08,  //
                                     0xA1, 0xA2, 0xA3, 0xA4,  //
                                     0xA5, 0xA6, 0xA7, 0xA8};

  iree_hal_buffer_t* device_buffer = NULL;
  IREE_ASSERT_OK(CreateZeroedDeviceBuffer(target_buffer_size, &device_buffer));

  iree_hal_command_buffer_t* command_buffer = NULL;
  IREE_ASSERT_OK(iree_hal_command_buffer_create(
      device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
      binding_capacity(), &command_buffer));
  IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));

  IREE_ASSERT_OK(iree_hal_command_buffer_update_buffer(
      command_buffer,
      /*source_buffer=*/source_buffer.data(), /*source_offset=*/0,
      MakeTargetRef(device_buffer, 0, target_buffer_size),
      IREE_HAL_UPDATE_FLAG_NONE));
  IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
  IREE_ASSERT_OK(SubmitWithBindings(command_buffer, device_buffer));

  std::vector<uint8_t> actual_data(target_buffer_size);
  IREE_ASSERT_OK(iree_hal_device_transfer_d2h(
      device_, device_buffer, /*source_offset=*/0, actual_data.data(),
      actual_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      iree_infinite_timeout()));
  EXPECT_THAT(actual_data, ContainerEq(source_buffer));

  iree_hal_command_buffer_release(command_buffer);
  iree_hal_buffer_release(device_buffer);
}

TEST_P(CommandBufferUpdateBufferTest, LargerPayload) {
  iree_device_size_t target_buffer_size = 256;
  std::vector<uint8_t> source_buffer(target_buffer_size);
  for (iree_host_size_t i = 0; i < source_buffer.size(); ++i) {
    source_buffer[i] = (uint8_t)(i ^ 0xA5u);
  }

  iree_hal_buffer_t* device_buffer = NULL;
  IREE_ASSERT_OK(CreateZeroedDeviceBuffer(target_buffer_size, &device_buffer));

  iree_hal_command_buffer_t* command_buffer = NULL;
  IREE_ASSERT_OK(iree_hal_command_buffer_create(
      device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
      binding_capacity(), &command_buffer));
  IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));

  IREE_ASSERT_OK(iree_hal_command_buffer_update_buffer(
      command_buffer,
      /*source_buffer=*/source_buffer.data(), /*source_offset=*/0,
      MakeTargetRef(device_buffer, 0, target_buffer_size),
      IREE_HAL_UPDATE_FLAG_NONE));
  IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
  IREE_ASSERT_OK(SubmitWithBindings(command_buffer, device_buffer));

  std::vector<uint8_t> actual_data(target_buffer_size);
  IREE_ASSERT_OK(iree_hal_device_transfer_d2h(
      device_, device_buffer, /*source_offset=*/0, actual_data.data(),
      actual_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      iree_infinite_timeout()));
  EXPECT_THAT(actual_data, ContainerEq(source_buffer));

  iree_hal_command_buffer_release(command_buffer);
  iree_hal_buffer_release(device_buffer);
}

TEST_P(CommandBufferUpdateBufferTest, WithOffsets) {
  iree_device_size_t target_buffer_size = 16;
  std::vector<uint8_t> source_buffer{0x01, 0x02, 0x03, 0x04,  //
                                     0x05, 0x06, 0x07, 0x08,  //
                                     0xA1, 0xA2, 0xA3, 0xA4,  //
                                     0xA5, 0xA6, 0xA7, 0xA8};

  iree_hal_buffer_t* device_buffer = NULL;
  IREE_ASSERT_OK(CreateZeroedDeviceBuffer(target_buffer_size, &device_buffer));

  iree_hal_command_buffer_t* command_buffer = NULL;
  IREE_ASSERT_OK(iree_hal_command_buffer_create(
      device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
      binding_capacity(), &command_buffer));
  IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));

  IREE_ASSERT_OK(iree_hal_command_buffer_update_buffer(
      command_buffer,
      /*source_buffer=*/source_buffer.data(), /*source_offset=*/4,
      MakeTargetRef(device_buffer, /*target_offset=*/4, /*length=*/8),
      IREE_HAL_UPDATE_FLAG_NONE));
  IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
  IREE_ASSERT_OK(SubmitWithBindings(command_buffer, device_buffer));

  std::vector<uint8_t> actual_data(target_buffer_size);
  IREE_ASSERT_OK(iree_hal_device_transfer_d2h(
      device_, device_buffer, /*source_offset=*/0, actual_data.data(),
      actual_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      iree_infinite_timeout()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x05, 0x06, 0x07, 0x08,  //
                                        0xA1, 0xA2, 0xA3, 0xA4,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  iree_hal_command_buffer_release(command_buffer);
  iree_hal_buffer_release(device_buffer);
}

TEST_P(CommandBufferUpdateBufferTest, Subspan) {
  iree_device_size_t target_buffer_size = 16;
  std::vector<uint8_t> source_buffer{0x01, 0x02, 0x03, 0x04,  //
                                     0x05, 0x06, 0x07, 0x08,  //
                                     0xA1, 0xA2, 0xA3, 0xA4,  //
                                     0xA5, 0xA6, 0xA7, 0xA8};

  iree_hal_buffer_t* device_buffer = NULL;
  IREE_ASSERT_OK(CreateZeroedDeviceBuffer(target_buffer_size, &device_buffer));

  // Create a subspan of the device buffer.
  iree_device_size_t subspan_length = 8;
  iree_hal_buffer_t* buffer_subspan = NULL;
  IREE_ASSERT_OK(
      iree_hal_buffer_subspan(device_buffer, /*byte_offset=*/4, subspan_length,
                              iree_allocator_system(), &buffer_subspan));

  iree_hal_command_buffer_t* command_buffer = NULL;
  IREE_ASSERT_OK(iree_hal_command_buffer_create(
      device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
      binding_capacity(), &command_buffer));
  IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));

  IREE_ASSERT_OK(iree_hal_command_buffer_update_buffer(
      command_buffer,
      /*source_buffer=*/source_buffer.data(), /*source_offset=*/4,
      MakeTargetRef(buffer_subspan, /*target_offset=*/4, /*length=*/4),
      IREE_HAL_UPDATE_FLAG_NONE));
  IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
  IREE_ASSERT_OK(SubmitWithBindings(command_buffer, buffer_subspan));

  // Check the full buffer contents.
  std::vector<uint8_t> actual_data(target_buffer_size);
  IREE_ASSERT_OK(iree_hal_device_transfer_d2h(
      device_, device_buffer, /*source_offset=*/0, actual_data.data(),
      actual_data.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      iree_infinite_timeout()));
  std::vector<uint8_t> reference_buffer{0x00, 0x00, 0x00, 0x00,  //
                                        0x00, 0x00, 0x00, 0x00,  //
                                        0x05, 0x06, 0x07, 0x08,  //
                                        0x00, 0x00, 0x00, 0x00};
  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));

  // Also check via the subspan.
  std::vector<uint8_t> actual_data_subspan(subspan_length);
  IREE_ASSERT_OK(iree_hal_device_transfer_d2h(
      device_, buffer_subspan, /*source_offset=*/0, actual_data_subspan.data(),
      actual_data_subspan.size(), IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT,
      iree_infinite_timeout()));
  std::vector<uint8_t> reference_buffer_subspan{0x00, 0x00, 0x00, 0x00,  //
                                                0x05, 0x06, 0x07, 0x08};
  EXPECT_THAT(actual_data_subspan, ContainerEq(reference_buffer_subspan));

  iree_hal_command_buffer_release(command_buffer);
  iree_hal_buffer_release(buffer_subspan);
  iree_hal_buffer_release(device_buffer);
}

CTS_REGISTER_COMMAND_BUFFER_TEST_SUITE(CommandBufferUpdateBufferTest);

}  // namespace iree::hal::cts
