//
// Created by Vicente Ferrari Smith on 13.02.26.
//

#include "renderer.h"
#include "../graphics.h"

#include <print>

#include "vulkan.h"
#include "../sprite.h"
#include <vma/vk_mem_alloc.h>
#include <slang/slang.h>

extern int32_t window_width;
extern int32_t window_height;

bool SortKey::operator<(const SortKey& b) const {
    if (depth != b.depth) return depth < b.depth;
    if (pipeline != b.pipeline) return pipeline < b.pipeline;
    return materialID < b.materialID;
}

Renderer::Renderer(GLFWwindow *window) {

    create_pipeline_layout();
    colored_quad_pipeline = create_graphics_pipeline<vertex_p2_s2_st2_col4_a1_u32>(
        device,
        pipelineLayout,
        swapchain_format.format,
        VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
        true
        );
    create_default_sampler();
    create_descriptor_pool();
    createFrameResources();
}

void Renderer::begin_frame() {
    commands.clear();
}

void Renderer::flush() {

}

void Renderer::submit_quad(glm::vec2 pos, glm::vec2 scale) {
    RenderCommand cmd {};
    cmd.pipeline = PipelineType::ColoredQuad;
    cmd.key = {
        (uint16_t) pos.y,
        0,
        (uint8_t) PipelineType::ColoredQuad
    };

    cmd.colored_quad = {
        .pos = pos,
        .scale = scale,
        .color = {0, 1, 1, 1},
    };

    commands.push_back(cmd);
}

void Renderer::submit_sprite(glm::vec2 pos, const sprite_t &sprite) {
    RenderCommand cmd {};
    cmd.pipeline = PipelineType::TexturedQuad;
    cmd.key = {
        (uint16_t) pos.y,
        0,
        (uint8_t) PipelineType::TexturedQuad
    };

    cmd.textured_quad = {
        .position = pos,
        .size = {0, 0},
        .uvMin = {0, 0},
        .uvMax = {0, 0},
        .color = {1, 1, 1, 1},
        .textureID = 0,
    };

    commands.push_back(cmd);

    // assert(started == true, "You can't submit without having started the renderer first.");
    // renderable : Renderable;
    // renderable.type = .Sprite;
    //
    // if sprite.window_space
    //     renderable.projection_type = .ORTHOGRAPHIC_WINDOW;
    // else
    //     renderable.projection_type = .ORTHOGRAPHIC_WORLD;
    //
    // renderable.pos                  = pos;
    // renderable.sprite.texture_sheet = sprite.texture_sheet;
    // renderable.sprite.texture_cell  = sprite.texture_cell;
    // renderable.sprite.origin        = sprite.origin;
    // renderable.sprite.scale         = sprite.scale;
    // renderable.sprite.colour        = sprite.colour;
    // renderable.sprite.alpha         = alpha;
    //
    // array_add(*renderer.renderable_list, renderable);
}

void Renderer::create_pipeline_layout() {
    std::array<VkDescriptorSetLayoutBinding, 1> bindings = {
        {
            {
                .binding = 0,
                .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                .descriptorCount = nextTextureSlot,
                .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
            }
        }
    };

    VkDescriptorBindingFlags flags[1] = {
        VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT
    };

    VkDescriptorSetLayoutBindingFlagsCreateInfo layoutFlags{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
        .bindingCount = 1,
        .pBindingFlags = flags
    };

    VkDescriptorSetLayoutCreateInfo dslci{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = &layoutFlags,
        // .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
        .bindingCount = bindings.size(),
        .pBindings = bindings.data()
    };

    vkCreateDescriptorSetLayout(device, &dslci, nullptr, &descriptor_set_layout);

    VkPushConstantRange push_constant{
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
        .offset = 0,
        .size = sizeof(glm::mat4),
    };

    VkPipelineLayoutCreateInfo plci{
        .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .setLayoutCount = 1,
        .pSetLayouts = &descriptor_set_layout,
        .pushConstantRangeCount = 1,
        .pPushConstantRanges = &push_constant,
    };

    vkCreatePipelineLayout(device, &plci, nullptr, &pipelineLayout);
}

void Renderer::create_frame_resources() {

    const VkSemaphoreCreateInfo seci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
    };

    VkFenceCreateInfo fenceInfo{
        .sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO,
        .flags = VK_FENCE_CREATE_SIGNALED_BIT,
    };

    VkCommandPoolCreateInfo cpci{
        .sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO,
        .flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT,
        .queueFamilyIndex = queueFamily
    };

    frames.resize(MAX_FRAMES_IN_FLIGHT);

    for (uint32_t i = 0; i < MAX_FRAMES_IN_FLIGHT; ++i) {
        Frame &frame = frames[i];

        vkCreateSemaphore(device, &seci, nullptr, &frame.imageAvailable);

        vkCreateFence(device, &fenceInfo, nullptr, &frame.in_flight_fence);


        vkCreateCommandPool(device, &cpci, nullptr, &frame.commandPool);

        const VkCommandBufferAllocateInfo cbai{
            .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
            .commandPool = frame.commandPool,
            .level = VK_COMMAND_BUFFER_LEVEL_PRIMARY,
            .commandBufferCount = 1
        };

        vkAllocateCommandBuffers(device, &cbai, &frame.command_buffer);

        VkBufferCreateInfo bufferInfo = {
            .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
            .size = 1024 * 1024 * 4,
            .usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
        };

        VmaAllocationCreateInfo allocCreateInfo = {};
        allocCreateInfo.usage = VMA_MEMORY_USAGE_AUTO;
        allocCreateInfo.flags = VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT |
            VMA_ALLOCATION_CREATE_HOST_ACCESS_ALLOW_TRANSFER_INSTEAD_BIT |
            VMA_ALLOCATION_CREATE_MAPPED_BIT;

        vmaCreateBuffer(
            allocator,
            &bufferInfo,
            &allocCreateInfo,
            &frame.vertexBuffer.buffer,
            &frame.vertexBuffer.allocation,
            &frame.vertexBuffer.info);
    }
}

AllocatedBuffer Renderer::create_buffer(size_t allocSize, VkBufferUsageFlags usage, VmaMemoryUsage memoryUsage) {
    // allocate buffer
    VkBufferCreateInfo bufferInfo = {.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO};
    bufferInfo.pNext = nullptr;
    bufferInfo.size = allocSize;

    bufferInfo.usage = usage;

    VmaAllocationCreateInfo vmaallocInfo = {};
    vmaallocInfo.usage = memoryUsage;
    vmaallocInfo.flags = VMA_ALLOCATION_CREATE_MAPPED_BIT;
    AllocatedBuffer newBuffer{};

    // allocate the buffer
    vmaCreateBuffer(allocator, &bufferInfo, &vmaallocInfo, &newBuffer.buffer, &newBuffer.allocation, &newBuffer.info);

    return newBuffer;
}

void Renderer::destroy_buffer(const AllocatedBuffer& buffer) {
    vmaDestroyBuffer(allocator, buffer.buffer, buffer.allocation);
}

// GPUMeshBuffers Renderer::uploadMesh(std::span<uint32_t> indices, std::span<vertex_p2_st2_col4_a1_u32> vertices) {
//     const size_t vertexBufferSize = vertices.size() * sizeof(vertex_p2_st2_col4_a1_u32);
//     const size_t indexBufferSize = indices.size() * sizeof(uint32_t);
//
//     GPUMeshBuffers newSurface;
//
//     //create vertex buffer
//     newSurface.vertexBuffer = create_buffer(vertexBufferSize, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
//         VMA_MEMORY_USAGE_GPU_ONLY);
//
//     //find the adress of the vertex buffer
//     VkBufferDeviceAddressInfo deviceAdressInfo{ .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,.buffer = newSurface.vertexBuffer.buffer };
//     newSurface.vertexBufferAddress = vkGetBufferDeviceAddress(device, &deviceAdressInfo);
//
//     //create index buffer
//     newSurface.indexBuffer = create_buffer(indexBufferSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
//         VMA_MEMORY_USAGE_GPU_ONLY);
//
//     AllocatedBuffer staging = create_buffer(vertexBufferSize + indexBufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VMA_MEMORY_USAGE_CPU_ONLY);
//
//     void* data = staging.allocation->GetMappedData();
//
//     // copy vertex buffer
//     memcpy(data, vertices.data(), vertexBufferSize);
//     // copy index buffer
//     memcpy((char*)data + vertexBufferSize, indices.data(), indexBufferSize);
//
//     immediate_submit([&](VkCommandBuffer cmd) {
//         VkBufferCopy vertexCopy{ 0 };
//         vertexCopy.dstOffset = 0;
//         vertexCopy.srcOffset = 0;
//         vertexCopy.size = vertexBufferSize;
//
//         vkCmdCopyBuffer(cmd, staging.buffer, newSurface.vertexBuffer.buffer, 1, &vertexCopy);
//
//         VkBufferCopy indexCopy{ 0 };
//         indexCopy.dstOffset = 0;
//         indexCopy.srcOffset = vertexBufferSize;
//         indexCopy.size = indexBufferSize;
//
//         vkCmdCopyBuffer(cmd, staging.buffer, newSurface.indexBuffer.buffer, 1, &indexCopy);
//     });
//
//     destroy_buffer(staging);
//
//     return newSurface;
//
// }

VkPipeline Renderer::get_pipeline(PipelineType type) const {
    switch (type) {
    case PipelineType::TexturedQuad: return textured_quad_pipeline;
    case PipelineType::ColoredQuad:  return colored_quad_pipeline;
    case PipelineType::Line:         return line_pipeline;
    default:                         return {};
    }
}

// void Renderer::bind_material(VkCommandBuffer cmd, uint16_t materialID) {
//     // In a real app, you'd have an array/map: std::vector<VkDescriptorSet> textureSets;
//     VkDescriptorSet set = textureSets[materialID];
//
//     vkCmdBindDescriptorSets(
//         cmd,
//         VK_PIPELINE_BIND_POINT_GRAPHICS,
//         pipelineLayout, // Our shared layout
//         0,              // Starting at Set 0
//         1,              // Binding 1 set
//         &set,
//         0, nullptr
//     );
// }

void Renderer::create_descriptor_pool() {
    VkDescriptorPoolSize pool_sizes[] = {
        { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, nextTextureSlot },
    };

    VkDescriptorPoolCreateInfo pool_info{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
        .maxSets = 1,
        .poolSizeCount = 1,
        .pPoolSizes = pool_sizes
    };

    vkCreateDescriptorPool(device, &pool_info, nullptr, &descriptorPool);

    VkDescriptorSetAllocateInfo alloc_info{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        .descriptorPool = descriptorPool,
        .descriptorSetCount = 1,
        .pSetLayouts = &descriptor_set_layout
    };

    vkAllocateDescriptorSets(device, &alloc_info, &set);
}

void Renderer::update_bindless_slot(uint32_t slot, VkImageView view, VkSampler sampler) const {
    VkDescriptorImageInfo image_info{
        .sampler = sampler,
        .imageView = view,
        .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL
    };

    VkWriteDescriptorSet write{
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = set,
        .dstArrayElement = slot,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        .pImageInfo = &image_info
    };

    vkUpdateDescriptorSets(device, 1, &write, 0, nullptr);
}

void Renderer::immediate_submit(std::function<void(VkCommandBuffer)>&& func) const {
    VkCommandBufferAllocateInfo allocInfo{ .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
    allocInfo.commandPool = frames[currentFrame].commandPool; // Use a pool created with VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = 1;

    VkCommandBuffer cmd;
    vkAllocateCommandBuffers(device, &allocInfo, &cmd);

    VkCommandBufferBeginInfo beginInfo{ .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

    vkBeginCommandBuffer(cmd, &beginInfo);

    // Execute the code passed in the lambda
    func(cmd);

    vkEndCommandBuffer(cmd);

    VkSubmitInfo submit{ .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO };
    submit.commandBufferCount = 1;
    submit.pCommandBuffers = &cmd;

    // Submit and wait
    vkQueueSubmit(graphics_queue, 1, &submit, VK_NULL_HANDLE);
    vkQueueWaitIdle(graphics_queue);

    vkFreeCommandBuffers(device, frames[currentFrame].commandPool, 1, &cmd);
}

void Renderer::transition_image_layout(VkCommandBuffer cmd, VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout) const {
    VkImageMemoryBarrier2 barrier{ .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2 };

    barrier.oldLayout = oldLayout;
    barrier.newLayout = newLayout;
    barrier.image = image;
    barrier.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };

    // Simple synchronization: wait for all previous commands, and block all future ones
    // You can optimize these masks later, but this is safe for a 2D engine
    barrier.srcStageMask = VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT;
    barrier.srcAccessMask = VK_ACCESS_2_MEMORY_WRITE_BIT;
    barrier.dstStageMask = VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT;
    barrier.dstAccessMask = VK_ACCESS_2_MEMORY_READ_BIT | VK_ACCESS_2_MEMORY_WRITE_BIT;

    VkDependencyInfo dep{ .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO };
    dep.imageMemoryBarrierCount = 1;
    dep.pImageMemoryBarriers = &barrier;

    vkCmdPipelineBarrier2(cmd, &dep);
}

VkImageView Renderer::create_image_view(VkImage image, VkFormat format) const {
    VkImageViewCreateInfo viewInfo{
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .image = image,
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .format = format,
    };

    // Default component mapping (R,G,B,A)
    viewInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
    viewInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
    viewInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
    viewInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;

    // Which part of the image to look at (Mip 0, Layer 0)
    viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    viewInfo.subresourceRange.baseMipLevel = 0;
    viewInfo.subresourceRange.levelCount = 1;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount = 1;

    VkImageView view;
    vkCreateImageView(device, &viewInfo, nullptr, &view);
    return view;
}

void Renderer::create_default_sampler() {
    VkSamplerCreateInfo samplerInfo{
        .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        // For crisp pixel art, use NEAREST. For smooth textures, use LINEAR.
        .magFilter = VK_FILTER_NEAREST,
        .minFilter = VK_FILTER_NEAREST,

        .mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST,

        // How to handle "out of bounds" UVs
        .addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,
        .addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE,

        // Optimization: turn off things we don't need for simple 2D
        .anisotropyEnable = VK_FALSE,
        .maxAnisotropy = 1.0f,
        .compareEnable = VK_FALSE,
        .compareOp = VK_COMPARE_OP_ALWAYS,
        .borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK,
        .unnormalizedCoordinates = VK_FALSE,
    };

    vkCreateSampler(device, &samplerInfo, nullptr, &defaultSampler);
}

void Renderer::end_frame() {

    Frame &frame = frames[currentFrame];

    vkWaitForFences(device, 1, &frame.in_flight_fence, VK_TRUE, UINT64_MAX);
    vkResetFences(device, 1, &frame.in_flight_fence);

    uint32_t imageIndex;
    vkAcquireNextImageKHR(
        device,
        swapchain,
        UINT64_MAX,
        frame.imageAvailable,
        VK_NULL_HANDLE,
        &imageIndex
    );

    commands = counting_sort_descending(commands, [](const RenderCommand &cmd){
        return cmd.key.depth;
    });

    std::vector<vertex_p2_s2_st2_col4_a1_u32> vertices;

    for (auto& cmd : commands) {


        switch (cmd.pipeline) {
        case PipelineType::ColoredQuad: {
            const auto &q = cmd.colored_quad;

            // Calculate spatial corners
            //float x0 = q.position.x;
            //float y0 = q.position.y;
            //float x1 = q.position.x + q.size.x;
            //float y1 = q.position.y + q.size.y;

            // Calculate UV corners
            // float u0 = q.uvMin.x;
            // float v0 = q.uvMin.y;
            // float u1 = q.uvMax.x;
            // float v1 = q.uvMax.y;

            // Define the 4 corners of the quad
            vertex_p2_s2_st2_col4_a1_u32 vTL = { q.pos, q.scale, {0, 0}, {1, 0, 0, 0}, 1, 0 };
            vertex_p2_s2_st2_col4_a1_u32 vTR = { q.pos, q.scale, {0, 0}, q.color, 1, 0 };
            vertex_p2_s2_st2_col4_a1_u32 vBL = { q.pos, q.scale, {0, 0}, q.color, 1, 0 };
            vertex_p2_s2_st2_col4_a1_u32 vBR = { q.pos, q.scale, {0, 0}, q.color, 1, 0 };

            // vertex_p2_st2_col4_a1_u32 vTL = { {x0, y0}, {u0, v0}, q.color, 1, q.textureID };
            // vertex_p2_st2_col4_a1_u32 vTR = { {x1, y0}, {u1, v0}, q.color, 1, q.textureID };
            // vertex_p2_st2_col4_a1_u32 vBL = { {x0, y1}, {u0, v1}, q.color, 1, q.textureID };
            // vertex_p2_st2_col4_a1_u32 vBR = { {x1, y1}, {u1, v1}, q.color, 1, q.textureID };

            vertices.push_back(vTL);
            vertices.push_back(vBL);
            vertices.push_back(vTR);

            vertices.push_back(vTR);
            vertices.push_back(vBL);
            vertices.push_back(vBR);

            break;
        }
        default:
            break;
        }
    }

    VkCommandBuffer cmd = frame.command_buffer;
    vkResetCommandBuffer(cmd, 0);

    VkCommandBufferBeginInfo cbBI {
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        .flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT
    };
    vkBeginCommandBuffer(cmd, &cbBI);

    recordCommandBuffer(
        cmd,
        images[imageIndex],
        imageViews[imageIndex],
        swapchain_extent,
        imageLayouts[imageIndex],
        frame,
        vertices
    );

    vkEndCommandBuffer(cmd);

    imageLayouts[imageIndex] = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

    VkSemaphoreSubmitInfo waitBinary{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_SUBMIT_INFO,
        .semaphore = frame.imageAvailable,
        .stageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT
    };

    VkSemaphoreSubmitInfo signalBinary{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_SUBMIT_INFO,
        .semaphore = renderFinished[imageIndex],
        .stageMask = VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT
    };

    VkCommandBufferSubmitInfo cmdInfo{
        .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_SUBMIT_INFO,
        .commandBuffer = cmd,
    };

    const VkSubmitInfo2 submit{
        .sType = VK_STRUCTURE_TYPE_SUBMIT_INFO_2,
        .waitSemaphoreInfoCount = 1,
        .pWaitSemaphoreInfos = &waitBinary,
        .commandBufferInfoCount = 1,
        .pCommandBufferInfos = &cmdInfo,
        .signalSemaphoreInfoCount = 1,
        .pSignalSemaphoreInfos = &signalBinary,
    };

    vkQueueSubmit2(graphics_queue, 1, &submit, frame.in_flight_fence);

    VkPresentInfoKHR present{
        .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
        .waitSemaphoreCount = 1,
        .pWaitSemaphores = &renderFinished[imageIndex],
        .swapchainCount = 1,
        .pSwapchains = &swapchain,
        .pImageIndices = &imageIndex,
    };

    vkQueuePresentKHR(graphics_queue, &present);

    currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
}

void Renderer::upload_vertex_buffer(
    VkCommandBuffer cmd,
    const Frame &frame,
    std::span<const vertex_p2_s2_st2_col4_a1_u32> vertices) const
{
    VkMemoryPropertyFlags memPropFlags;
    vmaGetAllocationMemoryProperties(allocator, frame.vertexBuffer.allocation, &memPropFlags);

    if(memPropFlags & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT) {
        // The Allocation ended up in a mappable memory.
        // Calling vmaCopyMemoryToAllocation() does vmaMapMemory(), memcpy(), vmaUnmapMemory(), and vmaFlushAllocation().
        VkResult result = vmaCopyMemoryToAllocation(allocator, vertices.data(), frame.vertexBuffer.allocation, 0, vertices.size() * sizeof(vertex_p2_s2_st2_col4_a1_u32));
        // Check result...

        VkBufferMemoryBarrier bufMemBarrier = { VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER };
        bufMemBarrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT;
        bufMemBarrier.dstAccessMask = VK_ACCESS_UNIFORM_READ_BIT;
        bufMemBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        bufMemBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        bufMemBarrier.buffer = frame.vertexBuffer.buffer;
        bufMemBarrier.offset = 0;
        bufMemBarrier.size = VK_WHOLE_SIZE;

        // It's important to insert a buffer memory barrier here to ensure writing to the buffer has finished.
        vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_HOST_BIT, VK_PIPELINE_STAGE_VERTEX_SHADER_BIT,
            0, 0, nullptr, 1, &bufMemBarrier, 0, nullptr);
    }
}

void Renderer::recordCommandBuffer(
    VkCommandBuffer cmd,
    VkImage image,
    VkImageView imageView,
    VkExtent2D extent,
    VkImageLayout oldLayout,
    const Frame &frame,
    const std::vector<vertex_p2_s2_st2_col4_a1_u32> &vertices) const
{

    {
        VkImageMemoryBarrier2 toColor{ .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2 };
        toColor.srcStageMask  = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT;
        toColor.dstStageMask  = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT;
        toColor.dstAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT;
        toColor.oldLayout     = VK_IMAGE_LAYOUT_UNDEFINED;
        toColor.newLayout     = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        toColor.image         = image;
        toColor.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 };

        VkDependencyInfo dep{
            .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
            .imageMemoryBarrierCount = 1,
            .pImageMemoryBarriers = &toColor
        };
        vkCmdPipelineBarrier2(cmd, &dep);
    }

    VkClearValue clearColor = {{{0.1f, 0.1f, 0.2f, 1.0f}}};
    VkRenderingAttachmentInfo colorAttach{
        .sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO,
        .imageView = imageView,
        .imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
        .loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR,
        .storeOp = VK_ATTACHMENT_STORE_OP_STORE,
        .clearValue = clearColor
    };

    VkRenderingInfo ri{
        .sType = VK_STRUCTURE_TYPE_RENDERING_INFO,
        .renderArea = {{0,0}, extent},
        .layerCount = 1,
        .colorAttachmentCount = 1,
        .pColorAttachments = &colorAttach
    };

    upload_vertex_buffer(cmd, frame, vertices);

    vkCmdBeginRendering(cmd, &ri);

    VkViewport vp{0.0f, 0.0f, (float)extent.width, (float)extent.height, 0.0f, 1.0f};
    VkRect2D sc{{0, 0}, extent};
    vkCmdSetViewport(cmd, 0, 1, &vp);
    vkCmdSetScissor(cmd, 0, 1, &sc);

    vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout, 0, 1, &set, 0, nullptr);
    VkDeviceSize vOffset{ 0 };
    vkCmdBindVertexBuffers(cmd, 0, 1, &frame.vertexBuffer.buffer, &vOffset);

    glm::mat4 projection = glm::ortho(0.0f, (float)window_width, 0.0f, (float)window_height, -1.0f, 1.0f);

    vkCmdPushConstants(
        cmd,
        pipelineLayout,
        VK_SHADER_STAGE_VERTEX_BIT,
        0,
        sizeof(glm::mat4),
        &projection
    );

    PipelineType lastPipeline = PipelineType::None; // Track current state
    // uint32_t vertexOffset = currentFrame * MAX_VERTICES_PER_BATCH;
    uint32_t currentBatchVertices = 0;

    for (const auto & render_command : commands) {
        // Only switch pipelines if we have to
        if (render_command.pipeline != lastPipeline) {
            // If we were mid-batch, draw what we have before switching
            if (currentBatchVertices > 0) {
                vkCmdDraw(cmd, currentBatchVertices, 1, 0, 0);
                // vertexOffset += currentBatchVertices;
                currentBatchVertices = 0;
            }

            vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, get_pipeline(render_command.pipeline));
            lastPipeline = render_command.pipeline;
        }

        currentBatchVertices += 6;
    }

    // Draw the final batch
    if (currentBatchVertices > 0) {
        vkCmdDraw(cmd, currentBatchVertices, 1, 0, 0);
    }

    vkCmdEndRendering(cmd);

    // 3. Transition back to Present
    {
        VkImageMemoryBarrier2 toPresent{
            .sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER_2,
            .srcStageMask  = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
            .srcAccessMask = VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT,
            .dstStageMask  = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
            .dstAccessMask = 0,
            .oldLayout     = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
            .newLayout     = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR,
            .image         = image,
            .subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 },
        };

        VkDependencyInfo dep{
            .sType = VK_STRUCTURE_TYPE_DEPENDENCY_INFO,
            .imageMemoryBarrierCount = 1,
            .pImageMemoryBarriers = &toPresent
        };
        
        vkCmdPipelineBarrier2(cmd, &dep);
    }

}
