#include "BackendResoureManager.hpp"

#include "MayaFlux/Core/Backends/Graphics/Vulkan/VKCommandManager.hpp"
#include "MayaFlux/Core/Backends/Graphics/Vulkan/VKContext.hpp"
#include "MayaFlux/Core/Backends/Graphics/Vulkan/VKImage.hpp"

#include "MayaFlux/Buffers/VKBuffer.hpp"
#include "MayaFlux/Journal/Archivist.hpp"
#include "MayaFlux/Registry/Service/BufferService.hpp"

namespace MayaFlux::Core {

namespace {
    struct FencedSubmission {
        vk::CommandBuffer cmd;
        vk::Fence fence;
    };
}

BackendResourceManager::BackendResourceManager(VKContext& context, VKCommandManager& command_manager)
    : m_context(context)
    , m_command_manager(command_manager)
{
}

void BackendResourceManager::setup_backend_service(const std::shared_ptr<Registry::Service::BufferService>& buffer_service)
{
    buffer_service->initialize_buffer = [this](const std::shared_ptr<void>& vk_buf) -> void {
        auto buffer = std::static_pointer_cast<Buffers::VKBuffer>(vk_buf);
        this->initialize_buffer(buffer);
    };

    buffer_service->allocate_raw_buffer = [this](
                                              size_t size_bytes,
                                              uint32_t usage_flags,
                                              uint32_t memory_property_flags,
                                              bool host_visible,
                                              void*& out_buffer,
                                              void*& out_memory,
                                              void*& out_mapped_ptr) -> void {
        vk::Buffer buf;
        vk::DeviceMemory mem;
        void* mapped = nullptr;
        this->allocate_raw_buffer(
            size_bytes,
            vk::BufferUsageFlags(usage_flags),
            vk::MemoryPropertyFlags(memory_property_flags),
            host_visible,
            buf, mem, mapped);
        out_buffer = static_cast<VkBuffer>(buf);
        out_memory = static_cast<VkDeviceMemory>(mem);
        out_mapped_ptr = mapped;
    };

    buffer_service->destroy_buffer = [this](const std::shared_ptr<void>& vk_buf) {
        auto buffer = std::static_pointer_cast<Buffers::VKBuffer>(vk_buf);
        this->cleanup_buffer(buffer);
    };

    buffer_service->get_buffer_device_address = [this](const std::shared_ptr<void>& vk_buf) -> uint64_t {
        auto buffer = std::static_pointer_cast<Buffers::VKBuffer>(vk_buf);
        auto address = this->get_buffer_device_address(buffer);
        return static_cast<uint64_t>(address);
    };

    buffer_service->execute_immediate = [this](const std::function<void(void*)>& recorder) {
        this->execute_immediate_commands([recorder](vk::CommandBuffer cmd) {
            recorder(static_cast<void*>(cmd));
        });
    };

    buffer_service->record_deferred = [this](const std::function<void(void*)>& recorder) {
        this->record_deferred_commands([recorder](vk::CommandBuffer cmd) {
            recorder(static_cast<void*>(cmd));
        });
    };

    buffer_service->flush_range = [this](void* memory, size_t offset, size_t size) {
        vk::DeviceMemory mem(reinterpret_cast<VkDeviceMemory>(memory));
        vk::MappedMemoryRange range { mem, offset, size == 0 ? VK_WHOLE_SIZE : size };
        if (auto result = m_context.get_device().flushMappedMemoryRanges(1, &range); result != vk::Result::eSuccess) {
            MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Failed to flush mapped memory range: {}", vk::to_string(result));
        }
    };

    buffer_service->invalidate_range = [this](void* memory, size_t offset, size_t size) {
        vk::DeviceMemory mem(reinterpret_cast<VkDeviceMemory>(memory));
        vk::MappedMemoryRange range { mem, offset, size == 0 ? VK_WHOLE_SIZE : size };
        if (auto result = m_context.get_device().invalidateMappedMemoryRanges(1, &range); result != vk::Result::eSuccess) {
            MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Failed to invalidate mapped memory range: {}", vk::to_string(result));
        }
    };

    buffer_service->map_buffer = [this](void* memory, size_t offset, size_t size) -> void* {
        vk::DeviceMemory mem(reinterpret_cast<VkDeviceMemory>(memory));
        return m_context.get_device().mapMemory(mem, offset, size == 0 ? VK_WHOLE_SIZE : size);
    };

    buffer_service->unmap_buffer = [this](void* memory) {
        vk::DeviceMemory mem(reinterpret_cast<VkDeviceMemory>(memory));
        m_context.get_device().unmapMemory(mem);
    };

    buffer_service->copy_buffer = [this](void* src, void* dst, size_t size, size_t src_off, size_t dst_off) {
        vk::Buffer s(static_cast<VkBuffer>(src));
        vk::Buffer d(static_cast<VkBuffer>(dst));
        execute_immediate_commands([&](vk::CommandBuffer cmd) {
            vk::BufferCopy region { src_off, dst_off, size };
            cmd.copyBuffer(s, d, 1, &region);
        });
    };

    buffer_service->execute_fenced = [this](const std::function<void(void*)>& recorder)
        -> std::shared_ptr<void> {
        vk::CommandBuffer cmd = m_command_manager.begin_single_time_commands();

        recorder(static_cast<void*>(cmd));

        cmd.end();

        auto device = m_context.get_device();
        vk::FenceCreateInfo fence_info {};
        vk::Fence fence = device.createFence(fence_info);

        vk::SubmitInfo submit_info {};
        submit_info.commandBufferCount = 1;
        submit_info.pCommandBuffers = &cmd;

        if (auto result = m_context.get_graphics_queue().submit(1, &submit_info, fence);
            result != vk::Result::eSuccess) {
            MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "execute_fenced: queue submit failed: {}", vk::to_string(result));
            device.destroyFence(fence);
            m_command_manager.free_command_buffer(cmd);
            return nullptr;
        }

        auto handle = std::make_shared<FencedSubmission>();
        handle->cmd = cmd;
        handle->fence = fence;
        return handle;
    };

    buffer_service->wait_fenced = [this](const std::shared_ptr<void>& handle) {
        if (!handle)
            return;
        auto sub = std::static_pointer_cast<FencedSubmission>(handle);
        if (!sub->fence)
            return;

        auto device = m_context.get_device();
        if (auto result = device.waitForFences(1, &sub->fence, VK_TRUE, UINT64_MAX);
            result != vk::Result::eSuccess) {
            MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "wait_fenced: waitForFences failed: {}", vk::to_string(result));
        }
    };

    buffer_service->release_fenced = [this](const std::shared_ptr<void>& handle) {
        if (!handle)
            return;
        auto sub = std::static_pointer_cast<FencedSubmission>(handle);

        auto device = m_context.get_device();
        if (sub->fence) {
            device.destroyFence(sub->fence);
            sub->fence = nullptr;
        }
        if (sub->cmd) {
            m_command_manager.free_command_buffer(sub->cmd);
            sub->cmd = nullptr;
        }
    };

    buffer_service->copy_buffer_fenced = [this, buffer_service](void* src, void* dst, size_t size, size_t src_off, size_t dst_off)
        -> std::shared_ptr<void> {
        vk::Buffer s(static_cast<VkBuffer>(src));
        vk::Buffer d(static_cast<VkBuffer>(dst));
        return buffer_service->execute_fenced([&](void* cmd_ptr) {
            vk::CommandBuffer cmd(static_cast<VkCommandBuffer>(cmd_ptr));
            vk::BufferCopy region { src_off, dst_off, static_cast<vk::DeviceSize>(size) };
            cmd.copyBuffer(s, d, 1, &region);
        });
    };
}

void BackendResourceManager::allocate_raw_buffer(
    size_t size_bytes,
    vk::BufferUsageFlags usage,
    vk::MemoryPropertyFlags memory_properties,
    bool host_visible,
    vk::Buffer& out_buffer,
    vk::DeviceMemory& out_memory,
    void*& out_mapped_ptr)
{
    vk::BufferCreateInfo buffer_info {};
    buffer_info.size = size_bytes;
    buffer_info.usage = usage;
    buffer_info.sharingMode = vk::SharingMode::eExclusive;

    try {
        out_buffer = m_context.get_device().createBuffer(buffer_info);
    } catch (const vk::SystemError& e) {
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to create raw VkBuffer: " + std::string(e.what()));
    }

    vk::MemoryRequirements mem_requirements = m_context.get_device().getBufferMemoryRequirements(out_buffer);

    vk::MemoryAllocateInfo alloc_info;
    alloc_info.allocationSize = mem_requirements.size;
    alloc_info.memoryTypeIndex = find_memory_type(mem_requirements.memoryTypeBits, memory_properties);

    try {
        out_memory = m_context.get_device().allocateMemory(alloc_info);
    } catch (const vk::SystemError& e) {
        m_context.get_device().destroyBuffer(out_buffer);
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to allocate raw VkDeviceMemory: " + std::string(e.what()));
    }

    try {
        m_context.get_device().bindBufferMemory(out_buffer, out_memory, 0);
    } catch (const vk::SystemError& e) {
        m_context.get_device().freeMemory(out_memory);
        m_context.get_device().destroyBuffer(out_buffer);
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to bind raw buffer memory: " + std::string(e.what()));
    }

    out_mapped_ptr = nullptr;
    if (host_visible) {
        try {
            out_mapped_ptr = m_context.get_device().mapMemory(out_memory, 0, size_bytes);
        } catch (const vk::SystemError& e) {
            m_context.get_device().freeMemory(out_memory);
            m_context.get_device().destroyBuffer(out_buffer);
            error_rethrow(
                Journal::Component::Core,
                Journal::Context::GraphicsBackend,
                std::source_location::current(),
                "Failed to map raw buffer memory: " + std::string(e.what()));
        }
    }
}

void BackendResourceManager::initialize_buffer(const std::shared_ptr<Buffers::VKBuffer>& buffer)
{
    if (!buffer) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Attempted to initialize null VKBuffer");
        return;
    }
    if (buffer->is_initialized()) {
        MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "VKBuffer already initialized, skipping");
        return;
    }

    vk::Buffer vk_buffer;
    vk::DeviceMemory memory;
    void* mapped_ptr = nullptr;

    allocate_raw_buffer(
        buffer->get_size_bytes(),
        buffer->get_usage_flags(),
        vk::MemoryPropertyFlags(buffer->get_memory_properties()),
        buffer->is_host_visible(),
        vk_buffer, memory, mapped_ptr);

    auto& resources = buffer->get_buffer_resources();
    resources.buffer = vk_buffer;
    resources.memory = memory;
    resources.mapped_ptr = mapped_ptr;

    m_managed_buffers.push_back(buffer);

    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "VKBuffer initialized: {} bytes, modality: {}, VkBuffer: {:p}",
        buffer->get_size_bytes(),
        Kakshya::modality_to_string(buffer->get_modality()),
        (void*)buffer->get_buffer());
}

void BackendResourceManager::cleanup_buffer(const std::shared_ptr<Buffers::VKBuffer>& buffer)
{
    if (!buffer) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Attempted to cleanup null VKBuffer");
        return;
    }

    auto it = std::ranges::find(m_managed_buffers, buffer);
    if (it == m_managed_buffers.end()) {
        return;
    }

    auto& res = it->get()->get_buffer_resources();

    if (res.mapped_ptr) {
        m_context.get_device().unmapMemory(res.memory);
    }

    if (res.index_buffer) {
        m_context.get_device().destroyBuffer(res.index_buffer);
    }

    if (res.index_memory) {
        m_context.get_device().freeMemory(res.index_memory);
    }

    if (res.buffer) {
        m_context.get_device().destroyBuffer(res.buffer);
    }

    if (res.memory) {
        m_context.get_device().freeMemory(res.memory);
    }

    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "VKBuffer cleaned up: {:p}", static_cast<void*>(res.buffer));

    m_managed_buffers.erase(it);
}

void BackendResourceManager::flush_pending_buffer_operations()
{
    for (auto& buffer_wrapper : m_managed_buffers) {
        auto& resources = buffer_wrapper->get_buffer_resources();
        auto dirty_ranges = buffer_wrapper->get_and_clear_dirty_ranges();
        if (!dirty_ranges.empty()) {
            for (auto [offset, size] : dirty_ranges) {
                vk::MappedMemoryRange range;
                range.memory = resources.memory;
                range.offset = offset;
                range.size = size;
                if (auto result = m_context.get_device().flushMappedMemoryRanges(1, &range); result != vk::Result::eSuccess) {
                    MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                        "Failed to flush mapped memory range: {}", vk::to_string(result));
                }
            }
            MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Flushed {} dirty ranges for buffer {:p}", dirty_ranges.size(),
                (void*)buffer_wrapper->get_buffer());
        }

        auto invalid_ranges = buffer_wrapper->get_and_clear_invalid_ranges();
        if (!invalid_ranges.empty()) {
            for (auto [offset, size] : invalid_ranges) {
                vk::MappedMemoryRange range;
                range.memory = buffer_wrapper->get_buffer_resources().memory;
                range.offset = offset;
                range.size = size;
                if (auto result = m_context.get_device().invalidateMappedMemoryRanges(1, &range); result != vk::Result::eSuccess) {
                    MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
                        "Failed to invalidate mapped memory range: {}", vk::to_string(result));
                }
            }
            MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Invalidated {} ranges for buffer {:p}", invalid_ranges.size(),
                (void*)buffer_wrapper->get_buffer());
        }
    }
}

vk::DeviceAddress BackendResourceManager::get_buffer_device_address(
    const std::shared_ptr<Buffers::VKBuffer>& buffer) const
{
    if (!buffer || !buffer->is_initialized()) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "get_buffer_device_address: buffer not initialized");
        return 0;
    }

    vk::BufferDeviceAddressInfo info {};
    info.buffer = buffer->get_buffer();
    // return static_cast<uint64_t>(m_context.get_device().getBufferAddress(info));
    return m_context.get_device().getBufferAddress(info);
}

void BackendResourceManager::initialize_image(const std::shared_ptr<VKImage>& image)
{
    if (!image) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Attempted to initialize null VKImage");
        return;
    }

    if (image->is_initialized()) {
        MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "VKImage already initialized, skipping");
        return;
    }

    // ========================================================================
    // Step 1: Create VkImage
    // ========================================================================

    vk::ImageCreateInfo image_info {};

    switch (image->get_type()) {
    case VKImage::Type::TYPE_1D:
        image_info.imageType = vk::ImageType::e1D;
        break;
    case VKImage::Type::TYPE_2D:
    case VKImage::Type::TYPE_CUBE:
        image_info.imageType = vk::ImageType::e2D;
        break;
    case VKImage::Type::TYPE_3D:
        image_info.imageType = vk::ImageType::e3D;
        break;
    }

    image_info.extent.width = image->get_width();
    image_info.extent.height = image->get_height();
    image_info.extent.depth = image->get_depth();
    image_info.mipLevels = image->get_mip_levels();
    image_info.arrayLayers = image->get_array_layers();
    image_info.format = image->get_format();
    image_info.tiling = vk::ImageTiling::eOptimal;
    image_info.initialLayout = vk::ImageLayout::eUndefined;
    image_info.usage = image->get_usage_flags();
    image_info.sharingMode = vk::SharingMode::eExclusive;
    image_info.samples = vk::SampleCountFlagBits::e1; // No MSAA for now
    image_info.flags = (image->get_type() == VKImage::Type::TYPE_CUBE)
        ? vk::ImageCreateFlagBits::eCubeCompatible
        : vk::ImageCreateFlags {};

    vk::Image vk_image;
    try {
        vk_image = m_context.get_device().createImage(image_info);
    } catch (const vk::SystemError& e) {
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to create VkImage: " + std::string(e.what()));
    }

    // ========================================================================
    // Step 2: Allocate memory
    // ========================================================================

    vk::MemoryRequirements mem_requirements;
    mem_requirements = m_context.get_device().getImageMemoryRequirements(vk_image);

    vk::MemoryAllocateInfo alloc_info {};
    alloc_info.allocationSize = mem_requirements.size;
    alloc_info.memoryTypeIndex = find_memory_type(
        mem_requirements.memoryTypeBits,
        image->get_memory_properties());

    vk::DeviceMemory memory;
    try {
        memory = m_context.get_device().allocateMemory(alloc_info);
    } catch (const vk::SystemError& e) {
        m_context.get_device().destroyImage(vk_image);
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to allocate VkDeviceMemory for image: " + std::string(e.what()));
    }

    // ========================================================================
    // Step 3: Bind memory to image
    // ========================================================================

    try {
        m_context.get_device().bindImageMemory(vk_image, memory, 0);
    } catch (const vk::SystemError& e) {
        m_context.get_device().freeMemory(memory);
        m_context.get_device().destroyImage(vk_image);
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to bind memory to VkImage: " + std::string(e.what()));
    }

    // ========================================================================
    // Step 4: Create image view
    // ========================================================================

    vk::ImageViewCreateInfo view_info {};

    switch (image->get_type()) {
    case VKImage::Type::TYPE_1D:
        view_info.viewType = (image->get_array_layers() > 1)
            ? vk::ImageViewType::e1DArray
            : vk::ImageViewType::e1D;
        break;
    case VKImage::Type::TYPE_2D:
        view_info.viewType = (image->get_array_layers() > 1)
            ? vk::ImageViewType::e2DArray
            : vk::ImageViewType::e2D;
        break;
    case VKImage::Type::TYPE_3D:
        view_info.viewType = vk::ImageViewType::e3D;
        break;
    case VKImage::Type::TYPE_CUBE:
        view_info.viewType = vk::ImageViewType::eCube;
        break;
    }

    view_info.image = vk_image;
    view_info.format = image->get_format();
    view_info.subresourceRange.aspectMask = image->get_aspect_flags();
    view_info.subresourceRange.baseMipLevel = 0;
    view_info.subresourceRange.levelCount = image->get_mip_levels();
    view_info.subresourceRange.baseArrayLayer = 0;
    view_info.subresourceRange.layerCount = image->get_array_layers();

    view_info.components.r = vk::ComponentSwizzle::eIdentity;
    view_info.components.g = vk::ComponentSwizzle::eIdentity;
    view_info.components.b = vk::ComponentSwizzle::eIdentity;
    view_info.components.a = vk::ComponentSwizzle::eIdentity;

    vk::ImageView image_view;
    try {
        image_view = m_context.get_device().createImageView(view_info);
    } catch (const vk::SystemError& e) {
        m_context.get_device().freeMemory(memory);
        m_context.get_device().destroyImage(vk_image);
        error_rethrow(
            Journal::Component::Core,
            Journal::Context::GraphicsBackend,
            std::source_location::current(),
            "Failed to create VkImageView: " + std::string(e.what()));
    }

    // ========================================================================
    // Step 5: Store handles in VKImage
    // ========================================================================

    VKImageResources resources {};
    resources.image = vk_image;
    resources.image_view = image_view;
    resources.memory = memory;
    resources.sampler = nullptr;

    image->set_image_resources(resources);
    image->set_current_layout(vk::ImageLayout::eUndefined);

    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "VKImage initialized: {}x{}x{}, format: {}, {} mips, {} layers",
        image->get_width(), image->get_height(), image->get_depth(),
        vk::to_string(image->get_format()),
        image->get_mip_levels(), image->get_array_layers());
}

void BackendResourceManager::cleanup_image(const std::shared_ptr<VKImage>& image)
{
    if (!image || !image->is_initialized()) {
        return;
    }

    const auto& resources = image->get_image_resources();

    if (resources.image_view) {
        m_context.get_device().destroyImageView(resources.image_view);
    }

    if (resources.image) {
        m_context.get_device().destroyImage(resources.image);
    }

    if (resources.memory) {
        m_context.get_device().freeMemory(resources.memory);
    }

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "VKImage cleaned up");
}

void BackendResourceManager::transition_image_layout(
    vk::Image image,
    vk::ImageLayout old_layout,
    vk::ImageLayout new_layout,
    uint32_t mip_levels,
    uint32_t array_layers,
    vk::ImageAspectFlags aspect_flags)
{
    execute_immediate_commands([&](vk::CommandBuffer cmd) {
        vk::ImageMemoryBarrier barrier {};
        barrier.oldLayout = old_layout;
        barrier.newLayout = new_layout;
        barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.image = image;
        barrier.subresourceRange.aspectMask = aspect_flags;
        barrier.subresourceRange.baseMipLevel = 0;
        barrier.subresourceRange.levelCount = mip_levels;
        barrier.subresourceRange.baseArrayLayer = 0;
        barrier.subresourceRange.layerCount = array_layers;

        vk::PipelineStageFlags src_stage;
        vk::PipelineStageFlags dst_stage;

        if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eTransferDstOptimal) {
            barrier.srcAccessMask = vk::AccessFlags {};
            barrier.dstAccessMask = vk::AccessFlagBits::eTransferWrite;
            src_stage = vk::PipelineStageFlagBits::eTopOfPipe;
            dst_stage = vk::PipelineStageFlagBits::eTransfer;
        } else if (old_layout == vk::ImageLayout::eTransferDstOptimal && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal) {
            barrier.srcAccessMask = vk::AccessFlagBits::eTransferWrite;
            barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;
            src_stage = vk::PipelineStageFlagBits::eTransfer;
            dst_stage = vk::PipelineStageFlagBits::eFragmentShader;
        } else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal) {
            barrier.srcAccessMask = vk::AccessFlags {};
            barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;
            src_stage = vk::PipelineStageFlagBits::eTopOfPipe;
            dst_stage = vk::PipelineStageFlagBits::eFragmentShader;
        } else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eColorAttachmentOptimal) {
            barrier.srcAccessMask = vk::AccessFlags {};
            barrier.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
            src_stage = vk::PipelineStageFlagBits::eTopOfPipe;
            dst_stage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
        } else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eDepthStencilAttachmentOptimal) {
            barrier.srcAccessMask = vk::AccessFlags {};
            barrier.dstAccessMask = vk::AccessFlagBits::eDepthStencilAttachmentRead | vk::AccessFlagBits::eDepthStencilAttachmentWrite;
            src_stage = vk::PipelineStageFlagBits::eTopOfPipe;
            dst_stage = vk::PipelineStageFlagBits::eEarlyFragmentTests;
        } else if (old_layout == vk::ImageLayout::eUndefined && new_layout == vk::ImageLayout::eGeneral) {
            barrier.srcAccessMask = vk::AccessFlags {};
            barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite;
            src_stage = vk::PipelineStageFlagBits::eTopOfPipe;
            dst_stage = vk::PipelineStageFlagBits::eComputeShader;
        } else if (old_layout == vk::ImageLayout::eGeneral && new_layout == vk::ImageLayout::eShaderReadOnlyOptimal) {
            barrier.srcAccessMask = vk::AccessFlagBits::eShaderWrite;
            barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;
            src_stage = vk::PipelineStageFlagBits::eComputeShader;
            dst_stage = vk::PipelineStageFlagBits::eComputeShader;
        } else if (old_layout == vk::ImageLayout::eShaderReadOnlyOptimal && new_layout == vk::ImageLayout::eGeneral) {
            barrier.srcAccessMask = vk::AccessFlagBits::eShaderRead;
            barrier.dstAccessMask = vk::AccessFlagBits::eShaderWrite;
            src_stage = vk::PipelineStageFlagBits::eComputeShader;
            dst_stage = vk::PipelineStageFlagBits::eComputeShader;
        } else {
            barrier.srcAccessMask = vk::AccessFlagBits::eMemoryRead | vk::AccessFlagBits::eMemoryWrite;
            barrier.dstAccessMask = vk::AccessFlagBits::eMemoryRead | vk::AccessFlagBits::eMemoryWrite;
            src_stage = vk::PipelineStageFlagBits::eAllCommands;
            dst_stage = vk::PipelineStageFlagBits::eAllCommands;

            MF_WARN(Journal::Component::Core, Journal::Context::GraphicsBackend,
                "Using generic image layout transition");
        }

        cmd.pipelineBarrier(
            src_stage, dst_stage,
            vk::DependencyFlags {},
            0, nullptr, // Memory barriers
            0, nullptr, // Buffer barriers
            1, &barrier // Image barriers
        );
    });

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Image layout transitioned: {} -> {}",
        vk::to_string(old_layout), vk::to_string(new_layout));
}

void BackendResourceManager::upload_image_data(
    std::shared_ptr<VKImage> image,
    const void* data,
    size_t size)
{
    if (!image || !data) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Invalid parameters for upload_image_data");
        return;
    }

    auto staging = std::make_shared<Buffers::VKBuffer>(
        size,
        Buffers::VKBuffer::Usage::STAGING,
        Kakshya::DataModality::IMAGE_COLOR);

    initialize_buffer(staging);

    void* mapped = staging->get_mapped_ptr();
    if (!mapped) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to map staging buffer for image upload");
        cleanup_buffer(staging);
        return;
    }

    std::memcpy(mapped, data, size);
    staging->mark_dirty_range(0, size);

    auto& resources = staging->get_buffer_resources();
    vk::MappedMemoryRange range { resources.memory, 0, VK_WHOLE_SIZE };

    if (auto result = m_context.get_device().flushMappedMemoryRanges(1, &range); result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to flush mapped memory range: {}", vk::to_string(result));
    }

    execute_immediate_commands([&](vk::CommandBuffer cmd) {
        vk::ImageMemoryBarrier barrier {};
        barrier.oldLayout = image->get_current_layout();
        barrier.newLayout = vk::ImageLayout::eTransferDstOptimal;
        barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.image = image->get_image();
        barrier.subresourceRange.aspectMask = image->get_aspect_flags();
        barrier.subresourceRange.baseMipLevel = 0;
        barrier.subresourceRange.levelCount = image->get_mip_levels();
        barrier.subresourceRange.baseArrayLayer = 0;
        barrier.subresourceRange.layerCount = image->get_array_layers();
        barrier.srcAccessMask = vk::AccessFlagBits::eShaderRead;
        barrier.dstAccessMask = vk::AccessFlagBits::eTransferWrite;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eFragmentShader,
            vk::PipelineStageFlagBits::eTransfer,
            vk::DependencyFlags {},
            0, nullptr, 0, nullptr, 1, &barrier);

        vk::BufferImageCopy region {};
        region.bufferOffset = 0;
        region.bufferRowLength = 0;
        region.bufferImageHeight = 0;
        region.imageSubresource.aspectMask = image->get_aspect_flags();
        region.imageSubresource.mipLevel = 0;
        region.imageSubresource.baseArrayLayer = 0;
        region.imageSubresource.layerCount = image->get_array_layers();
        region.imageOffset = vk::Offset3D { 0, 0, 0 };
        region.imageExtent = vk::Extent3D {
            image->get_width(),
            image->get_height(),
            image->get_depth()
        };

        cmd.copyBufferToImage(
            staging->get_buffer(),
            image->get_image(),
            vk::ImageLayout::eTransferDstOptimal,
            1, &region);

        barrier.oldLayout = vk::ImageLayout::eTransferDstOptimal;
        barrier.newLayout = vk::ImageLayout::eShaderReadOnlyOptimal;
        barrier.srcAccessMask = vk::AccessFlagBits::eTransferWrite;
        barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eTransfer,
            vk::PipelineStageFlagBits::eFragmentShader,
            vk::DependencyFlags {},
            0, nullptr, 0, nullptr, 1, &barrier);
    });

    image->set_current_layout(vk::ImageLayout::eShaderReadOnlyOptimal);

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Uploaded {} bytes to image {}x{}",
        size, image->get_width(), image->get_height());
}

void BackendResourceManager::upload_image_data(
    std::shared_ptr<VKImage> image,
    const void* data,
    size_t size,
    const std::shared_ptr<Buffers::VKBuffer>& staging, bool deferred)
{
    if (!image || !data || !staging) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Invalid parameters for upload_image_data_with_staging");
        return;
    }

    void* mapped = staging->get_mapped_ptr();
    if (!mapped) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "upload_image_data_with_staging: staging buffer has no mapped pointer");
        return;
    }

    std::memcpy(mapped, data, size);
    staging->mark_dirty_range(0, size);

    auto& resources = staging->get_buffer_resources();
    vk::MappedMemoryRange range { resources.memory, 0, VK_WHOLE_SIZE };

    if (auto result = m_context.get_device().flushMappedMemoryRanges(1, &range);
        result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "upload_image_data_with_staging: flush failed: {}", vk::to_string(result));
    }

    auto record_command = [&](vk::CommandBuffer cmd) {
        vk::ImageMemoryBarrier barrier {};
        barrier.oldLayout = image->get_current_layout();
        barrier.newLayout = vk::ImageLayout::eTransferDstOptimal;
        barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.image = image->get_image();
        barrier.subresourceRange = {
            image->get_aspect_flags(), 0,
            image->get_mip_levels(), 0,
            image->get_array_layers()
        };
        barrier.srcAccessMask = vk::AccessFlagBits::eShaderRead;
        barrier.dstAccessMask = vk::AccessFlagBits::eTransferWrite;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eFragmentShader,
            vk::PipelineStageFlagBits::eTransfer,
            {}, 0, nullptr, 0, nullptr, 1, &barrier);

        vk::BufferImageCopy region {};
        region.imageSubresource.aspectMask = image->get_aspect_flags();
        region.imageSubresource.layerCount = image->get_array_layers();
        region.imageOffset = vk::Offset3D { 0, 0, 0 };
        region.imageExtent = vk::Extent3D {
            image->get_width(),
            image->get_height(),
            image->get_depth()
        };

        cmd.copyBufferToImage(
            staging->get_buffer(),
            image->get_image(),
            vk::ImageLayout::eTransferDstOptimal,
            1, &region);

        barrier.oldLayout = vk::ImageLayout::eTransferDstOptimal;
        barrier.newLayout = vk::ImageLayout::eShaderReadOnlyOptimal;
        barrier.srcAccessMask = vk::AccessFlagBits::eTransferWrite;
        barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eTransfer,
            vk::PipelineStageFlagBits::eFragmentShader,
            {}, 0, nullptr, 0, nullptr, 1, &barrier);
    };

    if (deferred) {
        record_deferred_commands(record_command);
    } else {
        execute_immediate_commands(record_command);
    }

    image->set_current_layout(vk::ImageLayout::eShaderReadOnlyOptimal);

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "upload_image_data_with_staging: {} bytes to image {}x{}",
        size, image->get_width(), image->get_height());
}

void BackendResourceManager::download_image_data(
    std::shared_ptr<VKImage> image,
    void* data,
    size_t size,
    const std::shared_ptr<Buffers::VKBuffer>& staging,
    bool deferred,
    vk::ImageLayout restore_layout,
    vk::PipelineStageFlags restore_stage)
{
    if (!image || !data) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "download_image_data: invalid parameters");
        return;
    }

    std::shared_ptr<Buffers::VKBuffer> local_staging;
    Buffers::VKBuffer* active_staging {};
    if (!staging) {
        local_staging = std::make_shared<Buffers::VKBuffer>(
            size,
            Buffers::VKBuffer::Usage::STAGING,
            Kakshya::DataModality::IMAGE_COLOR);
        initialize_buffer(local_staging);
        active_staging = local_staging.get();
    } else {
        active_staging = staging.get();
    }

    auto record_command = [&](vk::CommandBuffer cmd) {
        vk::ImageMemoryBarrier barrier {};
        barrier.oldLayout = image->get_current_layout();
        barrier.newLayout = vk::ImageLayout::eTransferSrcOptimal;
        barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
        barrier.image = image->get_image();
        barrier.subresourceRange.aspectMask = image->get_aspect_flags();
        barrier.subresourceRange.baseMipLevel = 0;
        barrier.subresourceRange.levelCount = image->get_mip_levels();
        barrier.subresourceRange.baseArrayLayer = 0;
        barrier.subresourceRange.layerCount = image->get_array_layers();
        barrier.srcAccessMask = vk::AccessFlagBits::eShaderWrite | vk::AccessFlagBits::eShaderRead;
        barrier.dstAccessMask = vk::AccessFlagBits::eTransferRead;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eFragmentShader | vk::PipelineStageFlagBits::eComputeShader,
            vk::PipelineStageFlagBits::eTransfer,
            vk::DependencyFlags {}, {}, {}, barrier);

        vk::BufferImageCopy region {};
        region.bufferOffset = 0;
        region.bufferRowLength = 0;
        region.bufferImageHeight = 0;
        region.imageSubresource.aspectMask = image->get_aspect_flags();
        region.imageSubresource.mipLevel = 0;
        region.imageSubresource.baseArrayLayer = 0;
        region.imageSubresource.layerCount = image->get_array_layers();
        region.imageOffset = vk::Offset3D { 0, 0, 0 };
        region.imageExtent = vk::Extent3D {
            image->get_width(),
            image->get_height(),
            image->get_depth()
        };

        cmd.copyImageToBuffer(
            image->get_image(),
            vk::ImageLayout::eTransferSrcOptimal,
            active_staging->get_buffer(),
            1, &region);

        barrier.oldLayout = vk::ImageLayout::eTransferSrcOptimal;
        barrier.newLayout = restore_layout;
        barrier.srcAccessMask = vk::AccessFlagBits::eTransferRead;
        barrier.dstAccessMask = vk::AccessFlagBits::eMemoryRead;

        cmd.pipelineBarrier(
            vk::PipelineStageFlagBits::eTransfer,
            restore_stage,
            vk::DependencyFlags {}, {}, {}, barrier);
    };

    if (deferred && staging) {
        record_deferred_commands(record_command);
        image->set_current_layout(restore_layout);
        return;
    }

    auto device = m_context.get_device();

    vk::CommandBuffer cmd = m_command_manager.begin_single_time_commands();
    record_command(cmd);
    cmd.end();

    vk::FenceCreateInfo fence_info {};
    vk::Fence fence = device.createFence(fence_info);

    vk::SubmitInfo submit_info {};
    submit_info.commandBufferCount = 1;
    submit_info.pCommandBuffers = &cmd;

    if (auto result = m_context.get_graphics_queue().submit(1, &submit_info, fence);
        result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "download_image_data: queue submit failed: {}", vk::to_string(result));
        device.destroyFence(fence);
        m_command_manager.free_command_buffer(cmd);
        if (local_staging)
            cleanup_buffer(local_staging);
        return;
    }

    image->set_current_layout(restore_layout);

    if (auto result = device.waitForFences(1, &fence, VK_TRUE, UINT64_MAX);
        result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "download_image_data: waitForFences failed: {}", vk::to_string(result));
    }

    vk::MappedMemoryRange range { active_staging->get_buffer_resources().memory, 0, VK_WHOLE_SIZE };
    if (auto result = device.invalidateMappedMemoryRanges(1, &range);
        result != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "download_image_data: invalidate failed: {}", vk::to_string(result));
    }

    void* mapped = active_staging->get_mapped_ptr();
    if (mapped) {
        std::memcpy(data, mapped, size);
    } else {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "download_image_data: staging buffer has no mapped pointer");
    }

    device.destroyFence(fence);
    m_command_manager.free_command_buffer(cmd);

    if (local_staging)
        cleanup_buffer(local_staging);

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "download_image_data: {} bytes from image {}x{}",
        size, image->get_width(), image->get_height());
}

vk::Sampler BackendResourceManager::create_sampler(
    vk::Filter filter,
    vk::SamplerAddressMode address_mode,
    float max_anisotropy)
{
    size_t hash = 0;
    auto hash_combine = [](size_t& seed, size_t value) {
        seed ^= value + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    };

    hash_combine(hash, static_cast<size_t>(filter));
    hash_combine(hash, static_cast<size_t>(address_mode));
    hash_combine(hash, std::hash<float> {}(max_anisotropy));

    auto it = m_sampler_cache.find(hash);
    if (it != m_sampler_cache.end()) {
        MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Reusing cached sampler (hash: 0x{:X})", hash);
        return it->second;
    }

    vk::SamplerCreateInfo sampler_info;
    sampler_info.magFilter = filter;
    sampler_info.minFilter = filter;
    sampler_info.mipmapMode = vk::SamplerMipmapMode::eLinear;
    sampler_info.addressModeU = address_mode;
    sampler_info.addressModeV = address_mode;
    sampler_info.addressModeW = address_mode;
    sampler_info.mipLodBias = 0.0F;
    sampler_info.anisotropyEnable = max_anisotropy > 0.0F;
    sampler_info.maxAnisotropy = max_anisotropy;
    sampler_info.compareEnable = VK_FALSE;
    sampler_info.compareOp = vk::CompareOp::eAlways;
    sampler_info.minLod = 0.0F;
    sampler_info.maxLod = VK_LOD_CLAMP_NONE;
    sampler_info.borderColor = vk::BorderColor::eFloatOpaqueBlack;
    sampler_info.unnormalizedCoordinates = VK_FALSE;

    vk::Sampler sampler;
    try {
        sampler = m_context.get_device().createSampler(sampler_info);
    } catch (const vk::SystemError& e) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to create sampler: {}", e.what());
        return nullptr;
    }

    m_sampler_cache[hash] = sampler;

    MF_INFO(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Created sampler (filter: {}, address: {}, anisotropy: {}, hash: 0x{:X})",
        vk::to_string(filter), vk::to_string(address_mode), max_anisotropy, hash);

    return sampler;
}

void BackendResourceManager::destroy_sampler(vk::Sampler sampler)
{
    if (!sampler) {
        return;
    }

    for (auto it = m_sampler_cache.begin(); it != m_sampler_cache.end(); ++it) {
        if (it->second == sampler) {
            m_sampler_cache.erase(it);
            break;
        }
    }

    m_context.get_device().destroySampler(sampler);

    MF_DEBUG(Journal::Component::Core, Journal::Context::GraphicsBackend,
        "Destroyed sampler");
}

uint32_t BackendResourceManager::find_memory_type(uint32_t type_filter, vk::MemoryPropertyFlags properties) const
{
    vk::PhysicalDeviceMemoryProperties mem_properties;
    mem_properties = m_context.get_physical_device().getMemoryProperties();

    for (uint32_t i = 0; i < mem_properties.memoryTypeCount; i++) {
        if ((type_filter & (1 << i)) && (mem_properties.memoryTypes[i].propertyFlags & properties) == properties) {
            return i;
        }
    }

    error<std::runtime_error>(
        Journal::Component::Core,
        Journal::Context::GraphicsBackend,
        std::source_location::current(),
        "Failed to find suitable memory type");

    return 0;
}

void BackendResourceManager::execute_immediate_commands(const std::function<void(vk::CommandBuffer)>& recorder)
{
    vk::CommandBuffer cmd = m_command_manager.begin_single_time_commands();
    recorder(cmd);
    m_command_manager.end_single_time_commands(cmd, m_context.get_graphics_queue());
}

void BackendResourceManager::record_deferred_commands(const std::function<void(vk::CommandBuffer)>& recorder)
{
    m_command_manager.record_deferred_commands(recorder);
}

vk::Semaphore BackendResourceManager::flush_deferred_commands()
{
    if (!m_command_manager.has_deferred_commands())
        return nullptr;

    vk::CommandBuffer cmd = m_command_manager.get_deferred_cmd();
    cmd.end();

    if (!m_deferred_semaphore) {
        vk::SemaphoreCreateInfo sem_info {};
        m_deferred_semaphore = m_context.get_device().createSemaphore(sem_info);
    }

    vk::SubmitInfo submit {};
    submit.commandBufferCount = 1;
    submit.pCommandBuffers = &cmd;
    submit.signalSemaphoreCount = 1;
    submit.pSignalSemaphores = &m_deferred_semaphore;

    if (m_context.get_graphics_queue().submit(1, &submit, nullptr) != vk::Result::eSuccess) {
        MF_ERROR(Journal::Component::Core, Journal::Context::GraphicsBackend,
            "Failed to submit deferred command buffer");
    }

    m_command_manager.reset_deferred();

    return m_deferred_semaphore;
}

void BackendResourceManager::cleanup()
{
    for (auto& [hash, sampler] : m_sampler_cache) {
        if (sampler) {
            m_context.get_device().destroySampler(sampler);
        }
    }
    m_sampler_cache.clear();

    if (m_deferred_semaphore) {
        m_context.get_device().destroySemaphore(m_deferred_semaphore);
        m_deferred_semaphore = nullptr;
    }

    for (auto& buffer : m_managed_buffers) {
        if (buffer && buffer->is_initialized()) {
            cleanup_buffer(buffer);
        }
    }
    m_managed_buffers.clear();
}

size_t BackendResourceManager::compute_sampler_hash(vk::Filter filter, vk::SamplerAddressMode address_mode, float max_anisotropy) const
{
    size_t hash = 0;
    auto hash_combine = [](size_t& seed, size_t value) {
        seed ^= value + 0x9e3779b9 + (seed << 6) + (seed >> 2);
    };

    hash_combine(hash, static_cast<size_t>(filter));
    hash_combine(hash, static_cast<size_t>(address_mode));
    hash_combine(hash, std::hash<float> {}(max_anisotropy));

    return hash;
}

}
