#include "GPUBuffer.h"
#include "GPU.h"

GPUBuffer::GPUBuffer(SDL_GPUBufferUsageFlags usage, Uint32 size)
{
    SDL_GPUBufferCreateInfo buf_info = {
        .usage = usage,
        .size = size,
        .props = {}
    };
    this->size = size;
    buffer = SDL_CreateGPUBuffer(GPU::GetDevice(), &buf_info);
}

GPUBuffer::GPUBuffer(SDL_GPUBufferCreateInfo* info)
{
    this->size = info->size;
    buffer = SDL_CreateGPUBuffer(GPU::GetDevice(), info);
}

GPUBuffer::~GPUBuffer()
{
    SDL_ReleaseGPUBuffer(GPU::GetDevice(), buffer);
}

void* GPUBuffer::Transfer(void* data)
{
    return Transfer(data, size);
}

void* GPUBuffer::Transfer(void* data, Uint32 transferSize)
{
    return Transfer(data, transferSize, 0);
}

void* GPUBuffer::Transfer(void* data, Uint32 transferSize, Uint32 bufferOffset)
{
    if (!buffer || !data || transferSize == 0 || bufferOffset > size || transferSize > size - bufferOffset) {
        SDL_Log("Invalid GPU buffer transfer (buffer=%p, data=%p, size=%u, capacity=%u)",
                (void*)buffer, data, transferSize, size);
        return data;
    }

    SDL_GPUTransferBufferCreateInfo transfer_info = {
            .usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
            .size = transferSize,
            .props = {}
    };
    SDL_GPUTransferBuffer* transfer = SDL_CreateGPUTransferBuffer(GPU::GetDevice(), &transfer_info);
    if (!transfer) {
        SDL_Log("Failed to create GPU transfer buffer: %s", SDL_GetError());
        return data;
    }
    void* mapped = SDL_MapGPUTransferBuffer(GPU::GetDevice(), transfer, false);
    if (!mapped) {
        SDL_Log("Failed to map GPU transfer buffer: %s", SDL_GetError());
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
        return data;
    }
    memcpy(mapped, data, transferSize);
    SDL_UnmapGPUTransferBuffer(GPU::GetDevice(), transfer);

    // Copy to GPU buffer
    SDL_GPUCommandBuffer* cmd = SDL_AcquireGPUCommandBuffer(GPU::GetDevice());
    if (!cmd) {
        SDL_Log("Failed to acquire GPU command buffer: %s", SDL_GetError());
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
        return data;
    }
    SDL_GPUCopyPass* copy_pass = SDL_BeginGPUCopyPass(cmd);
    if (!copy_pass) {
        SDL_Log("Failed to begin GPU copy pass: %s", SDL_GetError());
        SDL_CancelGPUCommandBuffer(cmd);
        SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);
        return data;
    }
    SDL_GPUTransferBufferLocation src = { .transfer_buffer = transfer, .offset = 0 };
    SDL_GPUBufferRegion dst = { .buffer = buffer, .offset = bufferOffset, .size = transferSize };
    SDL_UploadToGPUBuffer(copy_pass, &src, &dst, false);
    SDL_EndGPUCopyPass(copy_pass);
    SDL_SubmitGPUCommandBuffer(cmd);
    SDL_WaitForGPUIdle(GPU::GetDevice());
    SDL_ReleaseGPUTransferBuffer(GPU::GetDevice(), transfer);

    return data;
}

SDL_GPUBuffer* GPUBuffer::Get()
{
    return buffer;
}