#include "kernel/devices/block_device_utils.hpp"

#include "kapi/system.hpp"

#include "kernel/devices/block_device.hpp"
#include "kapi/devices/device.hpp"

#include <kstd/cstring>
#include <kstd/memory>
#include <kstd/vector>

#include <algorithm>
#include <cstddef>

namespace kernel::devices::block_device_utils
{

  using block_op = void (*)(size_t idx, size_t off, size_t len, size_t done, devices::block_device * device,
                            std::byte * scratch, void * buffer);

  auto process_blocks(kstd::shared_ptr<kapi::devices::device> const & device, size_t offset, size_t size, void * buffer,
                      block_op op) -> size_t
  {
    if (buffer == nullptr)
    {
      kapi::system::panic("[FILESYSTEM] device_file::process_blocks called with null buffer.");
    }

    if (size == 0)
    {
      return 0;
    }

    auto * block_dev = static_cast<devices::block_device *>(device.get());
    if (block_dev == nullptr)
    {
      kapi::system::panic("[FILESYSTEM] device_file: expected block_device.");
    }

    size_t const block_size = block_dev->block_size();
    size_t const capacity = block_dev->capacity();

    if (offset >= capacity)
      return 0;
    size_t const total_to_process = std::min(size, capacity - offset);

    kstd::vector<std::byte> scratch_buffer{block_size};
    auto processed = 0uz;

    while (processed < total_to_process)
    {
      size_t const absolute_offset = offset + processed;
      size_t const block_index = absolute_offset / block_size;
      size_t const in_block_offset = absolute_offset % block_size;
      size_t const chunk_size = std::min(total_to_process - processed, block_size - in_block_offset);

      op(block_index, in_block_offset, chunk_size, processed, block_dev, scratch_buffer.data(), buffer);

      processed += chunk_size;
    }

    return processed;
  }

  auto read(kstd::shared_ptr<kapi::devices::device> const & device, void * buffer, size_t offset, size_t size) -> size_t
  {
    return process_blocks(device, offset, size, buffer,
                          [](size_t idx, size_t off, size_t len, size_t done, devices::block_device * device,
                             std::byte * scratch, void * buffer) {
                            auto * out = static_cast<std::byte *>(buffer);
                            if (off == 0 && len == device->block_size())
                            {
                              device->read_block(idx, out + done);
                            }
                            else
                            {
                              device->read_block(idx, scratch);
                              kstd::libc::memcpy(out + done, scratch + off, len);
                            }
                          });
  }

  auto write(kstd::shared_ptr<kapi::devices::device> const & device, void const * buffer, size_t offset, size_t size)
      -> size_t
  {
    return process_blocks(device, offset, size, const_cast<void *>(buffer),
                          [](size_t idx, size_t off, size_t len, size_t done, devices::block_device * device,
                             std::byte * scratch, void * buffer) {
                            auto const * in = static_cast<std::byte const *>(buffer);
                            if (off == 0 && len == device->block_size())
                            {
                              device->write_block(idx, in + done);
                            }
                            else
                            {
                              device->read_block(idx, scratch);
                              kstd::libc::memcpy(scratch + off, in + done, len);
                              device->write_block(idx, scratch);
                            }
                          });
  }

}  // namespace kernel::devices::block_device_utils