//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

module;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM
export module Crafter.Graphics:ComputeShader;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :Device;
import :DescriptorHeapVulkan;

export namespace Crafter {
    // Tier 1: thin compute-pipeline wrapper. Owns one VkPipeline. Use this to
    // dispatch the four standard UI shaders and any user-authored compute
    // shader that follows the ui-shared.glsl contract.
    //
    // The pipeline follows the device's binding model
    // (Device::descriptorHeapSupported):
    //   heap -> created with VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT and
    //           no pipeline layout; the bindless heap supplies all bindings
    //           and push constants travel via vkCmdPushDataEXT.
    //   sets -> created from the shader's <name>.sets.spv build against the
    //           heap object's shared pipeline layout; the bindless descriptor
    //           set supplies the bindings and push constants travel via
    //           vkCmdPushConstants.
    class ComputeShader {
    public:
        VkPipeline pipeline = VK_NULL_HANDLE;
        // Sets mode only: the pipeline layout of the heap this shader was
        // loaded against, needed again by vkCmdPushConstants. Not owned — it
        // belongs to that DescriptorHeapVulkan. VK_NULL_HANDLE in heap mode.
        VkPipelineLayout layout = VK_NULL_HANDLE;

        // NVIDIA descriptor-heap AS-read workaround (issue #15 / #7): set by
        // Load when this shader ray-queries the TLAS through the descriptor
        // heap and was rewritten to read its device address from a push
        // constant. `workaroundTlasPushOffset` is the byte offset of that member
        // (after the caller's own push payload, or 0 if the shader had none).
        // Tracked per-shader — a global is clobbered by whichever shader was
        // patched last. Both inert (false/0) on every other driver.
        bool workaroundNeedsTlas = false;
        std::uint32_t workaroundTlasPushOffset = 0;

        ComputeShader() = default;
        ComputeShader(const ComputeShader&) = delete;
        ComputeShader& operator=(const ComputeShader&) = delete;
        ComputeShader(ComputeShader&& other) noexcept;
        ComputeShader& operator=(ComputeShader&& other) noexcept;
        ~ComputeShader();

        // Loads a SPIR-V compute shader from disk and creates its pipeline.
        //
        // `spvPath` always names the descriptor-heap build (<name>.spv). On a
        // device without the heap, the <name>.sets.spv build next to it is
        // loaded instead, so the caller never branches on the binding model;
        // a missing .sets.spv throws, naming the file.
        //
        // `heap` must already be Initialize()d and must be the heap assigned
        // to Window::descriptorHeap for every frame this shader is dispatched
        // in: in sets mode the pipeline is built against that heap's layout,
        // and Window::Render binds that heap's descriptor set.
        void Load(const std::filesystem::path& spvPath, DescriptorHeapVulkan& heap);

        // Bind, push constants (if any), dispatch. Caller computes group counts
        // and is responsible for any inter-dispatch barriers (UIRenderer::Dispatch
        // wraps this with the standard write-after-write barrier).
        //
        // pushBytes should be a multiple of 4 (Vulkan requires it for both
        // push mechanisms). On a device without the descriptor heap the
        // payload travels as ordinary push constants, so it must also fit
        // the device's maxPushConstantsSize (256 bytes on common drivers); a
        // larger one throws std::length_error. Keep push blocks within that
        // if the shader is meant to run on both binding models.
        //
        // tlasAddress is the NVIDIA descriptor-heap AS-read workaround hook
        // (issue #15 / #7): a shader that ray-queries the TLAS through the
        // descriptor heap is rewritten to read its device address from a push
        // constant, so the caller must supply the active frame's TLAS address
        // (RenderingElement3D::tlases[frameIdx].address) here. It is pushed at
        // the shader's workaroundTlasPushOffset only when the shader was
        // rewritten (workaroundNeedsTlas) — ignored otherwise and on every
        // other driver, so shaders that don't touch an AS pass nothing.
        void Dispatch(VkCommandBuffer cmd,
                      const void* push, std::uint32_t pushBytes,
                      std::uint32_t gx,
                      std::uint32_t gy = 1,
                      std::uint32_t gz = 1,
                      VkDeviceAddress tlasAddress = 0) const;
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM