#include "iwa/buffer.hpp"

#include "iwa/device.hpp"

namespace iwa
{
Buffer::Buffer(ObjectPtr<Device> owner, const BufferCreationArgs& args) : super_t(std::move(owner))
{
    mHandle = getOwner()->getVkHandle().createBuffer(vk::BufferCreateInfo
    {
        .flags = args.flags,
        .size = args.size,
        .usage = args.usage,
        .sharingMode = args.sharingMode,
        .queueFamilyIndexCount = static_cast<std::uint32_t>(args.queueFamilyIndices.size()),
        .pQueueFamilyIndices = args.queueFamilyIndices.data()
    });
}

Buffer::~Buffer() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyBuffer);
}

void Buffer::allocateMemory(HostVisible hostVisible, HostCoherent hostCoherent)
{
    const vk::MemoryRequirements memoryRequirements = getOwner()->getVkHandle().getBufferMemoryRequirements(mHandle);
    const vk::MemoryPropertyFlags memoryFlags = hostVisible ?
                                                (vk::MemoryPropertyFlagBits::eHostVisible | (hostCoherent ? vk::MemoryPropertyFlagBits::eHostCoherent : vk::MemoryPropertyFlags())) :
                                                vk::MemoryPropertyFlagBits::eDeviceLocal;
    const std::optional<std::uint32_t> memoryTypeIdx = findMemoryType(*getOwner(), memoryRequirements, memoryFlags);
    if (!memoryTypeIdx.has_value())
    {
        throw std::runtime_error("Could not find a suitable memory type.");
    }

    ObjectPtr<DeviceMemory> memory = getOwner()->allocateDeviceMemory(
    {
        .allocationSize = memoryRequirements.size,
        .memoryTypeIndex = memoryTypeIdx.value()
    });
    bindMemory(std::move(memory));
}

void Buffer::bindMemory(ObjectPtr<DeviceMemory> memory, vk::DeviceSize offset)
{
    mMemory = std::move(memory);
    getOwner()->getVkHandle().bindBufferMemory(mHandle, *mMemory, offset);
}

mijin::Task<> Buffer::c_fill(std::uint32_t data, std::size_t bytes, std::size_t byteOffset)
{
    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;
    cmdBuffer.fillBuffer(mHandle, byteOffset, bytes, data);

    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

mijin::Task<> Buffer::c_copyFrom(vk::Buffer srcBuffer, vk::BufferCopy region)
{
    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;
    cmdBuffer.copyBuffer(srcBuffer, mHandle, region);

    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

mijin::Task<> Buffer::c_upload(const void* data, std::size_t bytes, std::size_t byteOffset)
{
    // assert(bytes == SIZE_REST || bytes + byteOffset <= byteSize);
    // if (bytes == SIZE_REST) {
    //     bytes = byteSize - byteOffset;
    // }
    // create scratch buffer
    vk::Device device = *getOwner();
    ObjectPtr<Buffer> scratchBuffer = getOwner()->createChild<Buffer>(BufferCreationArgs{
        .size = bytes,
        .usage = vk::BufferUsageFlagBits::eTransferSrc
    });
    scratchBuffer->allocateMemory(HostVisible::YES);

    // copy to scratch buffer
    void* mapped = device.mapMemory(*scratchBuffer->getMemory(), 0, bytes);
    std::memcpy(mapped, data, bytes);
    device.unmapMemory(*scratchBuffer->getMemory());

    // copy to actual buffer
    co_await c_copyFrom(*scratchBuffer, vk::BufferCopy{
        .srcOffset = 0,
        .dstOffset = byteOffset,
        .size = bytes
    });
}

mijin::Task<> Buffer::c_upload(const mijin::TypelessBuffer& data, std::size_t byteOffset)
{
    return c_upload(data.data(), data.byteSize(), byteOffset);
}
} // namespace iwa