/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

// Issue #115: the frame loop's inter-pass and frame-edge (acquire/present)
// barriers used to set BOTH stage masks to VK_PIPELINE_STAGE_ALL_COMMANDS_BIT
// and the inter-pass dependency to a queue-wide VkMemoryBarrier — flushing
// every cache and fully serialising against the whole pipeline every frame,
// when all the next pass needs is the swapchain image the previous one wrote.
//
// The fix scopes the inter-pass dependency to the swapchain image's single
// colour subresource (BuildSwapchainInterPassBarrier) and narrows the stage
// masks to the stages that actually touch that image, derived per pass via the
// polymorphic RenderPass::SwapchainStage() and unioned across the frame by
// SwapchainStageUnion. The load-bearing correctness constraint is that this
// derivation is per-pass: a compute pass writes via COMPUTE_SHADER, a ray-
// tracing pass via RAY_TRACING_SHADER, so a hardcoded compute mask would
// UNDER-synchronise an RT pass and corrupt the image. All three helpers are
// pure CPU logic over the pass list, so this test drives them directly with no
// GPU device at runtime — mirroring MipChainBarrierBatch / UploadStrategy.

#include <cstdlib>
#include "vulkan/vulkan.h"

import Crafter.Graphics;
import std;
using namespace Crafter;

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

// A pass that never overrides SwapchainStage(): stands in for any future pass
// that forgets to narrow its stage, and must fall back to the conservative
// writer union rather than silently picking one stage (which could under-sync).
struct DefaultPass : RenderPass {
    void Record(GraphicsCommandBuffer, std::uint32_t, Window&) override {}
};

VkImage SentinelImage() {
    return reinterpret_cast<VkImage>(static_cast<std::uintptr_t>(0xC0FFEE));
}

} // namespace

int main() {
    // ─── per-pass SwapchainStage() polymorphism ─────────────────────────────
    // The whole reason the fix is "moderate, not trivial": these two stages
    // MUST differ, or scoping an RT pass's barrier to the compute stage would
    // under-synchronise it. Pin that they are distinct and each correct.
    Check(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT != VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
          "compute and ray-tracing write stages are distinct (the under-sync trap)");

    UIRenderer ui;                 // real compute UI pass
    RTPass rt(nullptr);            // real RT pass (SwapchainStage doesn't touch the pipeline)
    DefaultPass def;               // un-overridden fallback

    // Access through the base pointer — the frame loop only ever sees RenderPass*.
    RenderPass* uiBase  = &ui;
    RenderPass* rtBase  = &rt;
    RenderPass* defBase = &def;

    Check(uiBase->SwapchainStage() == VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
          "UIRenderer reports the COMPUTE_SHADER stage");
    Check(rtBase->SwapchainStage() == VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
          "RTPass reports the RAY_TRACING_SHADER stage");
    Check(defBase->SwapchainStage() == kSwapchainWriterStages,
          "an un-overridden pass falls back to the conservative writer union");

    // ─── the conservative writer union constant ─────────────────────────────
    constexpr VkPipelineStageFlags expectedUnion =
        VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT
      | VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR
      | VK_PIPELINE_STAGE_TRANSFER_BIT;
    Check(kSwapchainWriterStages == expectedUnion,
          "writer union = COMPUTE_SHADER | RAY_TRACING_SHADER | TRANSFER");
    // The point of the whole change: even the conservative fallback is far
    // tighter than the ALL_COMMANDS mask it replaces.
    Check(kSwapchainWriterStages != VK_PIPELINE_STAGE_ALL_COMMANDS_BIT,
          "writer union is narrower than ALL_COMMANDS");

    // ─── SwapchainStageUnion over the frame's pass list ─────────────────────
    // Empty frame: no writers to derive from, so fall back to the conservative
    // union (the image is still transitioned, just never written).
    Check(SwapchainStageUnion(std::span<RenderPass* const>{}) == kSwapchainWriterStages,
          "empty pass list falls back to the conservative writer union");

    // All-compute frame: the union narrows to COMPUTE_SHADER ONLY — this is the
    // real perf win, no RAY_TRACING / TRANSFER / ALL_COMMANDS dragged in.
    {
        std::array<RenderPass*, 2> computeOnly = {uiBase, uiBase};
        VkPipelineStageFlags u = SwapchainStageUnion(computeOnly);
        Check(u == VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
              "all-compute frame narrows the union to COMPUTE_SHADER only");
        Check((u & VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR) == 0,
              "all-compute frame does NOT pull in RAY_TRACING_SHADER");
    }

    // RT-only frame: narrows to RAY_TRACING_SHADER only.
    {
        std::array<RenderPass*, 1> rtOnly = {rtBase};
        Check(SwapchainStageUnion(rtOnly) == VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR,
              "RT-only frame narrows the union to RAY_TRACING_SHADER only");
    }

    // Mixed frame: the union folds in BOTH stages — never under-syncs the RT
    // pass, never over-syncs to ALL_COMMANDS.
    {
        std::array<RenderPass*, 2> mixed = {uiBase, rtBase};
        VkPipelineStageFlags u = SwapchainStageUnion(mixed);
        Check(u == (VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR),
              "mixed compute+RT frame unions COMPUTE_SHADER | RAY_TRACING_SHADER");
        Check((u & VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR) != 0,
              "mixed frame keeps RAY_TRACING_SHADER (no RT under-sync)");
    }

    // A frame containing an un-overridden pass conservatively widens to the
    // full writer union — safe by construction.
    {
        std::array<RenderPass*, 2> withDefault = {uiBase, defBase};
        Check(SwapchainStageUnion(withDefault) == kSwapchainWriterStages,
              "a frame with an un-overridden pass widens to the conservative union");
    }

    // ─── BuildSwapchainInterPassBarrier ─────────────────────────────────────
    // Replaces the queue-wide VkMemoryBarrier: an image memory barrier scoped
    // to the swapchain image's single colour subresource, GENERAL->GENERAL (no
    // layout change — the storage image stays bound), matching the intra-pass
    // UI barrier's access scope.
    {
        VkImageMemoryBarrier b = BuildSwapchainInterPassBarrier(SentinelImage());
        Check(b.sType == VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER, "inter-pass barrier sType is IMAGE_MEMORY_BARRIER");
        Check(b.image == SentinelImage(), "inter-pass barrier carries the swapchain image handle");
        Check(b.oldLayout == VK_IMAGE_LAYOUT_GENERAL && b.newLayout == VK_IMAGE_LAYOUT_GENERAL,
              "inter-pass barrier is GENERAL->GENERAL (pure memory dependency, no transition)");
        Check(b.srcAccessMask == VK_ACCESS_SHADER_WRITE_BIT,
              "inter-pass barrier src access is shader write");
        Check(b.dstAccessMask == (VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT),
              "inter-pass barrier dst access is shader read|write");
        Check(b.srcQueueFamilyIndex == VK_QUEUE_FAMILY_IGNORED && b.dstQueueFamilyIndex == VK_QUEUE_FAMILY_IGNORED,
              "inter-pass barrier is not a queue-family transfer");
        Check(b.subresourceRange.aspectMask == VK_IMAGE_ASPECT_COLOR_BIT,
              "inter-pass barrier covers the colour aspect");
        Check(b.subresourceRange.levelCount == 1 && b.subresourceRange.layerCount == 1,
              "inter-pass barrier is scoped to the single swapchain subresource");
    }

    if (failures != 0) {
        std::println("{} check(s) failed", failures);
        return EXIT_FAILURE;
    }
    std::println("all checks passed");
    return EXIT_SUCCESS;
}
