// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkWebGPUInternalsBuffer.h"
#include "vtkArrayDispatch.h"
#include "vtkLogger.h"
#include "vtkWGPUContext.h"

namespace
{
class DispatchDataWriter
{
public:
  DispatchDataWriter(wgpu::Device device, wgpu::Buffer buffer, vtkIdType byteOffset)
    : Device(device)
    , Buffer(buffer)
    , ByteOffset(byteOffset)
  {
  }

  template <typename SrcArrayType>
  void operator()(SrcArrayType* srcArray)
  {
    using SrcType = vtk::GetAPIType<SrcArrayType>;

    const auto srcValuesIterator = vtk::DataArrayValueRange(srcArray);

    std::vector<SrcType> data;
    data.reserve(srcValuesIterator.size());
    for (const auto& value : srcValuesIterator)
    {
      data.push_back(value);
    }

    this->Device.GetQueue().WriteBuffer(
      this->Buffer, this->ByteOffset, data.data(), data.size() * srcArray->GetDataTypeSize());
  }

private:
  wgpu::Device Device;
  wgpu::Buffer Buffer;
  vtkIdType ByteOffset;
};
}

//------------------------------------------------------------------------------
wgpu::Buffer vtkWebGPUInternalsBuffer::Upload(const wgpu::Device& device, unsigned long offset,
  void* data, unsigned long sizeBytes, wgpu::BufferUsage usage, const char* label /*=nullptr*/)
{
  wgpu::BufferDescriptor descriptor;
  descriptor.label = label == nullptr ? "(nolabel)" : label;
  descriptor.size = sizeBytes;
  descriptor.usage = usage | wgpu::BufferUsage::CopyDst;

  wgpu::Buffer buffer = device.CreateBuffer(&descriptor);
  device.GetQueue().WriteBuffer(buffer, offset, data, sizeBytes);
  return buffer;
}

//------------------------------------------------------------------------------
void vtkWebGPUInternalsBuffer::UploadFromDataArray(
  wgpu::Device device, wgpu::Buffer buffer, vtkDataArray* dataArray)
{
  UploadFromDataArray(device, buffer, 0, dataArray);
}

//------------------------------------------------------------------------------
void vtkWebGPUInternalsBuffer::UploadFromDataArray(
  wgpu::Device device, wgpu::Buffer buffer, vtkIdType byteOffset, vtkDataArray* dataArray)
{
  using DispatchAllTypes = vtkArrayDispatch::DispatchByValueType<vtkArrayDispatch::AllTypes>;

  DispatchDataWriter dispatchDataWriter(device, buffer, byteOffset);

  if (!DispatchAllTypes::Execute(dataArray, dispatchDataWriter))
  {
    dispatchDataWriter(dataArray);
  }
}

//------------------------------------------------------------------------------
wgpu::Buffer vtkWebGPUInternalsBuffer::CreateABuffer(const wgpu::Device& device,
  unsigned long sizeBytes, wgpu::BufferUsage usage, bool mappedAtCreation /*=false*/,
  const char* label /*=nullptr*/)
{
  if (!vtkWebGPUInternalsBuffer::CheckBufferSize(device, sizeBytes))
  {
    wgpu::SupportedLimits supportedDeviceLimits;
    device.GetLimits(&supportedDeviceLimits);

    vtkLog(ERROR,
      "The current WebGPU Device cannot create buffers larger than: "
        << supportedDeviceLimits.limits.maxStorageBufferBindingSize
        << " bytes but the buffer with label " << (label ? label : "") << " is " << sizeBytes
        << " bytes big.");

    return nullptr;
  }

  wgpu::BufferDescriptor descriptor;
  descriptor.label = label == nullptr ? "(nolabel)" : label;
  descriptor.size = sizeBytes;
  descriptor.usage = usage;
  descriptor.mappedAtCreation = mappedAtCreation;

  wgpu::Buffer buffer = device.CreateBuffer(&descriptor);
  return buffer;
}

//------------------------------------------------------------------------------
bool vtkWebGPUInternalsBuffer::CheckBufferSize(const wgpu::Device& device, unsigned long sizeBytes)
{
  wgpu::SupportedLimits supportedDeviceLimits;
  device.GetLimits(&supportedDeviceLimits);

  if (sizeBytes > supportedDeviceLimits.limits.maxStorageBufferBindingSize)
  {
    return false;
  }

  return true;
}
