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

#include "renderer.h"
#include "init.h"
#include "sprite.h"
#include "swapchain.h"

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) {

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

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

    create_pipeline_layout();
    colored_quad_pipeline = create_graphics_pipeline<vertex_p2_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() {
    uint64_t waitValue = 0;

    if (frameValue >= MAX_FRAMES_IN_FLIGHT) {
        waitValue = frameValue - MAX_FRAMES_IN_FLIGHT + 1;

        VkSemaphoreWaitInfo waitInfo{
            .sType = VK_STRUCTURE_TYPE_SEMAPHORE_WAIT_INFO,
            .semaphoreCount = 1,
            .pSemaphores = &timelineSemaphore,
            .pValues = &waitValue
        };

        vkWaitSemaphores(device, &waitInfo, UINT64_MAX);
    }
}

void Renderer::flush() {

}

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

    cmd.colored_quad = {
        .position = pos,
        .size = {0.25, 0.25},
        .color = {1, 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() {
    VkDescriptorSetLayoutBinding bindings[2];

    bindings[0] = VkDescriptorSetLayoutBinding{
        .binding = 0,
        .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        .descriptorCount = 1000,
        .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
    };

    bindings[1] = {
        .binding = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT // The vertex shader "pulls" from here
    };

    VkDescriptorBindingFlags flags[2] = {
        VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT,
        0 // Standard behavior for the buffer
    };

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

    VkDescriptorSetLayoutCreateInfo dslci{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = &layoutFlags, // Attach the flags
        .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
        .bindingCount = 2,
        .pBindings = bindings
    };

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

    // 2. Create the Shared Pipeline Layout
    VkPushConstantRange push_constant{
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
        .offset = 0,
        .size = sizeof(glm::mat4) // Camera Projection
    };

    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::createFrameResources() {
    VkSemaphoreTypeCreateInfo typeInfo{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_TYPE_CREATE_INFO,
        .semaphoreType = VK_SEMAPHORE_TYPE_TIMELINE,
        .initialValue = 0
    };

    const VkSemaphoreCreateInfo semaphoreci{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO,
        .pNext = &typeInfo
    };

    vkCreateSemaphore(device, &semaphoreci, nullptr, &timelineSemaphore);

    const VkSemaphoreCreateInfo seci{ VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO };

    for (auto & frame : frames) {
        vkCreateSemaphore(device, &seci, nullptr, &frame.imageAvailable);
        vkCreateSemaphore(device, &seci, nullptr, &frame.renderFinished);
    }

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

    for (auto &frame : frames) {
        vkAllocateCommandBuffers(device, &cbai, &frame.command_buffer);
    }

    VkBufferCreateInfo bufferInfo = { VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
    bufferInfo.size = 1024 * 1024 * 4;
    bufferInfo.usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;

    VmaAllocationCreateInfo allocInfo = {};
    allocInfo.usage = VMA_MEMORY_USAGE_CPU_TO_GPU;
    allocInfo.flags = VMA_ALLOCATION_CREATE_MAPPED_BIT;

    vmaCreateBuffer(allocator, &bufferInfo, &allocInfo, &vertexBuffer, &vertexAllocation, &vertexAllocInfo);

    VkDescriptorBufferInfo vertexBufferInfo{
        .buffer = vertexBuffer,
        .offset = 0,
        .range  = VK_WHOLE_SIZE
    };

    VkWriteDescriptorSet bufferWrite{
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = set,
        .dstBinding = 1,
        .dstArrayElement = 0,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .pBufferInfo = &vertexBufferInfo
    };

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

void Renderer::end_frame() {

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

    vertex_p2_st2_col4_a1_u32 *vPtr = (vertex_p2_st2_col4_a1_u32 *) vertexAllocInfo.pMappedData;
    vertex_p2_st2_col4_a1_u32 *currentFrameStart = vPtr + (currentFrame * MAX_VERTICES_PER_FRAME);

    uint32_t totalVertices = 0;
    for (auto& cmd : commands) {

        vPtr = currentFrameStart + totalVertices;

        switch (cmd.pipeline) {
        case PipelineType::ColoredQuad: {
            const auto &q = cmd.textured_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_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 };

            // --- Triangle 1 (TL, TR, BL) ---
            vPtr[0] = vTL;
            vPtr[1] = vTR;
            vPtr[2] = vBL;

            // --- Triangle 2 (TR, BR, BL) ---
            vPtr[3] = vTR;
            vPtr[4] = vBR;
            vPtr[5] = vBL;

            break;
        }
        default:
            break;
        }

        totalVertices += 6;
    }

    uint32_t imageIndex;
    vkAcquireNextImageKHR(
        device,
        swapchain,
        UINT64_MAX,
        frames[currentFrame].imageAvailable,
        VK_NULL_HANDLE,
        &imageIndex
    );

    VkCommandBuffer command_buffer = frames[currentFrame].command_buffer;
    vkResetCommandBuffer(command_buffer, 0);

    recordCommandBuffer(
        command_buffer,
        images[imageIndex],
        imageViews[imageIndex],
        swapchain_extent,
        imageLayouts[imageIndex]
    );

    commands.clear();

    imageLayouts[imageIndex] = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;

    frameValue++;

    VkSemaphoreSubmitInfo waitBinary{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_SUBMIT_INFO,
        .semaphore = frames[currentFrame].imageAvailable,
        .stageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT
    };

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

    VkSemaphoreSubmitInfo signalTimeline{
        .sType = VK_STRUCTURE_TYPE_SEMAPHORE_SUBMIT_INFO,
        .semaphore = timelineSemaphore,
        .value = frameValue,
        .stageMask = VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT
    };

    VkSemaphoreSubmitInfo signals[] = { signalBinary, signalTimeline };

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

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

    vkQueueSubmit2(graphics_queue, 1, &submit, VK_NULL_HANDLE);

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

    vkQueuePresentKHR(graphics_queue, &present);

    currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
}

void Renderer::recordCommandBuffer(
    VkCommandBuffer command_buffer,
    VkImage image,
    VkImageView imageView,
    VkExtent2D extent,
    VkImageLayout oldLayout) const
{
    VkCommandBufferBeginInfo begin{ .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
    vkBeginCommandBuffer(command_buffer, &begin);

    {
        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     = oldLayout;
        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(command_buffer, &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
    };

    vkCmdBeginRendering(command_buffer, &ri);

    vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout, 0, 1, &set, 0, nullptr);

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

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

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

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

        currentBatchVertices += 6;
    }

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

    vkCmdEndRendering(command_buffer);

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

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

    vkEndCommandBuffer(command_buffer);
}

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, 1000 },
        { VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1 },
    };

    VkDescriptorPoolCreateInfo pool_info{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
        .maxSets = 1000,
        .poolSizeCount = 2,
        .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
    };

    if (vkAllocateDescriptorSets(device, &alloc_info, &set) != VK_SUCCESS) {
        throw std::runtime_error("Failed to allocate bindless descriptor set!");
    }
}

void Renderer::update_bindless_slot(uint32_t slot, VkImageView view, VkSampler sampler) {
    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, // Index in the 1000-size array
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        .pImageInfo = &image_info
    };

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

Texture Renderer::upload_texture(int w, int h, void* pixels) {
    VkDeviceSize imageSize = w * h * 4;
    Texture res{};

    // --- 1. Create Staging Buffer (CPU Visible) ---
    VkBuffer stagingBuffer;
    VmaAllocation stagingAlloc;

    VkBufferCreateInfo stagingBufferInfo = { .sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO };
    stagingBufferInfo.size = imageSize;
    stagingBufferInfo.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;

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

    VmaAllocationInfo stagingResultInfo;
    vmaCreateBuffer(allocator, &stagingBufferInfo, &stagingAllocCreateInfo, &stagingBuffer, &stagingAlloc, &stagingResultInfo);

    // Copy raw pixels into the mapped memory provided by VMA
    memcpy(stagingResultInfo.pMappedData, pixels, imageSize);

    // --- 2. Create GPU Image (Device Local / Tiled) ---
    VkExtent3D imageExtent = { (uint32_t) w, (uint32_t) h, 1 };

    VkImageCreateInfo imageInfo = { .sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO };
    imageInfo.imageType = VK_IMAGE_TYPE_2D;
    imageInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    imageInfo.extent = imageExtent;
    imageInfo.mipLevels = 1;
    imageInfo.arrayLayers = 1;
    imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
    imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;

    VmaAllocationCreateInfo imageAllocCreateInfo = {};
    imageAllocCreateInfo.usage = VMA_MEMORY_USAGE_AUTO;
    imageAllocCreateInfo.priority = 1.0f; // High priority for textures

    vmaCreateImage(allocator, &imageInfo, &imageAllocCreateInfo, &res.image, &res.allocation, nullptr);

    // --- 3. The Transfer ---
    immediate_submit([&](VkCommandBuffer cmd) {
        // Transition image from UNDEFINED to TRANSFER_DST
        transition_image_layout(cmd, res.image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

        VkBufferImageCopy copyRegion = {};
        copyRegion.bufferOffset = 0;
        copyRegion.bufferRowLength = 0;
        copyRegion.bufferImageHeight = 0;
        copyRegion.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
        copyRegion.imageSubresource.mipLevel = 0;
        copyRegion.imageSubresource.baseArrayLayer = 0;
        copyRegion.imageSubresource.layerCount = 1;
        copyRegion.imageExtent = imageExtent;

        vkCmdCopyBufferToImage(cmd, stagingBuffer, res.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &copyRegion);

        // Transition image from TRANSFER_DST to SHADER_READ_ONLY
        transition_image_layout(cmd, res.image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    });

    // Clean up temporary staging resources
    vmaDestroyBuffer(allocator, stagingBuffer, stagingAlloc);

    // --- 4. Finalize Handles ---
    res.view = create_image_view(res.image, imageInfo.format);

    // Register in your Bindless Array (Set 0, Binding 0, Index N)
    res.descriptor_index = nextTextureSlot++;
    update_bindless_slot(res.descriptor_index, res.view, defaultSampler);

    return res;
}

void Renderer::immediate_submit(std::function<void(VkCommandBuffer)>&& func) const {
    VkCommandBufferAllocateInfo allocInfo{ .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
    allocInfo.commandPool = 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, 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;
    if (vkCreateImageView(device, &viewInfo, nullptr, &view) != VK_SUCCESS) {
        throw std::runtime_error("failed to create image 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.
    samplerInfo.magFilter = VK_FILTER_NEAREST;
    samplerInfo.minFilter = VK_FILTER_NEAREST;

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

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

    if (vkCreateSampler(device, &samplerInfo, nullptr, &defaultSampler) != VK_SUCCESS) {
        throw std::runtime_error("failed to create texture sampler!");
    }
}