//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:RTPass;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :RenderPass;
import :Window;
import :Device;
import :PipelineRTVulkan;
import :RenderingElement3D;

export namespace Crafter {
    struct RTPass : RenderPass {
        PipelineRTVulkan* pipeline;

        RTPass(PipelineRTVulkan* p) : pipeline(p) {}

        // An RT pass writes the swapchain image from the ray-tracing pipeline,
        // so the frame loop's barriers must wait on RAY_TRACING_SHADER — using
        // the compute stage here would under-synchronise and corrupt the image.
        VkPipelineStageFlags SwapchainStage() const override {
            return VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
        }

        void Record(VkCommandBuffer cmd, std::uint32_t frameIdx, Window& window) override {
            vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, pipeline->pipeline);
            // NVIDIA descriptor-heap AS-read workaround (issue #15 / #7): feed
            // the active frame's TLAS device address into the push-constant
            // block that VulkanShader synthesizes, so the rewritten raygen can
            // reach the acceleration structure by address instead of through
            // the faulting heap descriptor. Inert on every other driver.
            if (Device::workaroundDescriptorHeapAS && pipeline->workaroundNeedsTlas) {
                VkDeviceAddress tlasAddr = RenderingElement3D::tlases[frameIdx].address;
                VkPushDataInfoEXT pushInfo {
                    .sType  = VK_STRUCTURE_TYPE_PUSH_DATA_INFO_EXT,
                    // Where the rewritten raygen reads the TLAS address: 0 when
                    // VulkanShader synthesized a fresh block, or the offset of
                    // the member it appended to the shader's existing block.
                    // Tracked per-pipeline (copied from the shader table) so a
                    // later-loaded shader can't clobber it.
                    .offset = pipeline->workaroundTlasPushOffset,
                    .data   = { .address = &tlasAddr, .size = sizeof(tlasAddr) },
                };
                Device::vkCmdPushDataEXT(cmd, &pushInfo);
            }
            Device::vkCmdTraceRaysKHR(cmd,
                &pipeline->raygenRegion,
                &pipeline->missRegion,
                &pipeline->hitRegion,
                &pipeline->callableRegion,
                window.width, window.height, 1);
        }
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM

#ifdef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :RenderPass;
import :Window;
import :WebGPU;
import :PipelineRTWebGPU;
import :RenderingElement3D;

export namespace Crafter {
    // DOM-mode RT pass — dispatches the megakernel pipeline at frame Record
    // time. Picks up the current TLAS for the frame and the application's
    // raygen-side push data (typically empty in v1; pass via window.passes
    // wiring if needed later).
    struct RTPass : RenderPass {
        PipelineRTWebGPU* pipeline;
        // Optional per-dispatch push data forwarded after the standard
        // RTDispatchHeader. Null means "no extra data".
        const void*       pushPtr   = nullptr;
        std::uint32_t     pushBytes = 0;
        // Resolved WebGPU resource handles for each user binding the
        // pipeline was loaded with, in declaration order. The example
        // owns the storage (typically a small std::array of u32). Null /
        // 0 means "no user bindings".
        const void*       handlesPtr   = nullptr;
        std::uint32_t     handlesCount = 0;
        // Wavefront bounce budget: number of (TRACE; SHADE) iterations.
        // 1 = primary rays only; 2 = primary + one continuation/shadow
        // bounce; etc. The library unrolls GENERATE; (PREP; TRACE; SHADE)
        // ×maxDepth; RESOLVE.
        std::uint32_t     maxDepth     = 1;
        // Destination storage-texture handle for an HDR-output pipeline
        // (one Init'd with a non-default hdrOutputFormat). RESOLVE writes
        // the linear accumulator here instead of the canvas; the app runs
        // its own composite→swapchain pass afterwards. Ignored (0) for the
        // default canvas path.
        std::uint32_t     outTexHandle = 0;
        // Per-bounce ray budget as a multiple of the pixel count. The
        // wavefront ray/hit/payload buffers hold raysPerPixel·W·H rays, so
        // a closest-hit may emit up to raysPerPixel rays per pixel within
        // one bounce (e.g. one shadow ray per light — rtAccumulate is
        // atomic, issue #30) before rtEmitRay starts dropping. Memory
        // scales linearly; keep at 1 for single-ray-per-pixel pipelines.
        std::uint32_t     raysPerPixel = 1;

        RTPass(PipelineRTWebGPU* p) : pipeline(p) {}

        void Record(WebGPUCommandEncoderRef /*cmd*/, std::uint32_t frameIdx, Window& window) override {
            const std::uint32_t gx = (window.width  + 7u) / 8u;
            const std::uint32_t gy = (window.height + 7u) / 8u;
            auto& tlas = RenderingElement3D::tlases[frameIdx];
            WebGPU::wgpuDispatchRT(
                pipeline->pipelineHandle,
                pushPtr,
                static_cast<std::int32_t>(pushBytes),
                tlas.buffer.handle,
                static_cast<std::int32_t>(tlas.builtInstanceCount),
                static_cast<std::int32_t>(gx),
                static_cast<std::int32_t>(gy),
                handlesPtr,
                static_cast<std::int32_t>(handlesCount),
                static_cast<std::int32_t>(maxDepth),
                outTexHandle,
                static_cast<std::int32_t>(raysPerPixel));
        }
    };
}
#endif // CRAFTER_GRAPHICS_WINDOW_DOM