// Copyright 2025 Steven Le Rouzic
//
// SPDX-License-Identifier: BSD-3-Clause

#include "hk21/vulkan/sync/sync.hpp"

#include <asl/containers/buffer.hpp>
#include <asl/testing/testing.hpp>

class DependencyInfoBuilder : public vulkan_sync::DependencyInfoBuilder
{
    asl::buffer<VkImageMemoryBarrier2>  m_image_barriers;
    asl::buffer<VkBufferMemoryBarrier2> m_buffer_barriers;

public:
    void add_image_barrier(const VkImageMemoryBarrier2& barrier) override
    {
        m_image_barriers.push(barrier);
    }

    void add_buffer_barrier(const VkBufferMemoryBarrier2& barrier) override
    {
        m_buffer_barriers.push(barrier);
    }

    void reset()
    {
        m_image_barriers.clear();
        m_buffer_barriers.clear();
    }

    asl::span<const VkImageMemoryBarrier2> image_barriers() const
    {
        return m_image_barriers;
    }

    asl::span<const VkBufferMemoryBarrier2> buffer_barriers() const
    {
        return m_buffer_barriers;
    }
};

ASL_TEST(clear_and_present)
{
    DependencyInfoBuilder builder;
    vulkan_sync::ImageState state{};

    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageClear, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    auto barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_NONE);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageClear, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImagePresent, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_NONE);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageClear, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_NONE);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
}

ASL_TEST(clear_and_draw)
{
    DependencyInfoBuilder builder;
    vulkan_sync::ImageState state{};

    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageColorWriteAttachment, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    auto barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_BOTTOM_OF_PIPE_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_NONE);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageSampledInVertexShader, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_COLOR_ATTACHMENT_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_SHADER_SAMPLED_READ_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageSampledInFragmentShader, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask ==  VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT);
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_SHADER_SAMPLED_READ_BIT);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_SHADER_SAMPLED_READ_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_UNDEFINED);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageSampledInVertexShader, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 0);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    builder.reset();
    synchronize_resource(VK_NULL_HANDLE, {}, &state, vulkan_sync::Usage::kImageClear, &builder);
    ASL_TEST_ASSERT(builder.buffer_barriers().size() == 0);
    ASL_TEST_ASSERT(builder.image_barriers().size() == 1);
    barrier = builder.image_barriers()[0];
    ASL_TEST_EXPECT(barrier.srcStageMask == (VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT | VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT));
    ASL_TEST_EXPECT(barrier.srcAccessMask == VK_ACCESS_2_SHADER_SAMPLED_READ_BIT);
    ASL_TEST_EXPECT(barrier.dstStageMask == VK_PIPELINE_STAGE_2_CLEAR_BIT);
    ASL_TEST_EXPECT(barrier.dstAccessMask == VK_ACCESS_2_TRANSFER_WRITE_BIT);
    ASL_TEST_EXPECT(barrier.oldLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
    ASL_TEST_EXPECT(barrier.newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
    ASL_TEST_EXPECT(state.current_layout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
}