#pragma once
#include "DeferredGPUBuffer.h"
#include "WickedEngine.h"

namespace voxel {

// ── Ray Tracing Manager (Phase 6) ──────────────────────────────
// Groups all RT state: BLAS/TLAS management, shadow/AO dispatches.
// Extracted from VoxelRenderer to isolate the ~500 lines of RT code
// and its 20+ members for easier debugging and maintenance.

class VoxelRTManager {
public:
    // ── Initialization ──────────────────────────────────────────
    void initialize(wi::graphics::GraphicsDevice* device, uint32_t maxBlasVertices);

    // ── BLAS extraction (compute shaders) ───────────────────────

    // Extract blocky quad positions into BLAS vertex buffer.
    void dispatchBLASExtract(wi::graphics::CommandList cmd,
        const wi::graphics::GPUBuffer& quadBuffer,
        const wi::graphics::GPUBuffer& chunkInfoBuffer,
        uint32_t quadCount) const;

    // Extract toping instance positions via GPU compute.
    // groupBuffer/groupsGPU: toping BLAS group table.
    void dispatchTopingBLASExtract(wi::graphics::CommandList cmd,
        const wi::graphics::GPUBuffer& topingVertexBuffer,
        const wi::graphics::GPUBuffer& topingInstanceBuffer,
        const void* groupsGPUData, size_t groupsGPUSize,
        uint32_t groupCount, uint32_t totalVertices) const;

    // ── Acceleration structure build ────────────────────────────
    static constexpr uint32_t BUILD_BLOCKY = 1 << 0;
    static constexpr uint32_t BUILD_SMOOTH = 1 << 1;
    static constexpr uint32_t BUILD_TOPING = 1 << 2;
    static constexpr uint32_t BUILD_ALL    = BUILD_BLOCKY | BUILD_SMOOTH | BUILD_TOPING;

    void buildAccelerationStructures(wi::graphics::CommandList cmd,
        uint32_t buildFlags,
        const wi::graphics::GPUBuffer& smoothVB,
        uint32_t smoothVertCount) const;

    // ── RT Shadows + AO dispatch ────────────────────────────────
    void dispatchShadows(wi::graphics::CommandList cmd,
        const wi::graphics::Texture& depthBuffer,
        const wi::graphics::Texture& renderTarget,
        const wi::graphics::Texture& normalTarget,
        const wi::graphics::GPUBuffer& constantBuffer) const;

    // ── Toping BLAS buffer management ───────────────────────────
    // Ensure capacity for toping BLAS position + index buffers.
    // Returns true if buffers were (re)created.
    bool ensureTopingBLASCapacity(uint32_t totalVertices);

    // ── State queries ───────────────────────────────────────────
    bool isAvailable() const { return available_; }
    bool isReady() const { return available_ && tlas_.IsValid(); }
    bool isShadowsEnabled() const { return shadowsEnabled_; }
    void setShadowsEnabled(bool v) { shadowsEnabled_ = v; }
    uint32_t getShadowDebug() const { return shadowDebug_; }
    void setShadowDebug(uint32_t v) { shadowDebug_ = v; }

    uint32_t getBlockyTriCount() const { return blockyVertexCount_ / 3; }
    uint32_t getSmoothTriCount() const { return smoothVertexCount_ / 3; }
    uint32_t getTopingTriCount() const { return topingVertexCount_ / 3; }
    uint32_t getTopingVertexCount() const { return topingVertexCount_; }
    uint32_t getTlasInstanceCount() const { return tlasInstanceCount_; }
    const wi::graphics::RaytracingAccelerationStructure& getTLAS() const { return tlas_; }

    // Dirty flags (public for VoxelRenderPath orchestration)
    mutable bool dirty = true;           // BLAS/TLAS need rebuild
    mutable bool topingBLASDirty = false; // toping BLAS extract + rebuild needed
    mutable bool aoHistoryValid = false;
    mutable uint32_t frameCounter = 0;
    mutable XMFLOAT4X4 prevViewProjection;

    // AO textures (created by VoxelRenderPath::createRenderTargets)
    mutable wi::graphics::Texture aoRawTexture;
    mutable wi::graphics::Texture aoBlurredTexture;
    mutable wi::graphics::Texture aoHistoryTexture;

private:
    wi::graphics::GraphicsDevice* device_ = nullptr;
    mutable bool available_ = false;
    mutable bool shadowsEnabled_ = false;
    mutable uint32_t shadowDebug_ = 0;

    // Shaders
    wi::graphics::Shader blasExtractShader_;
    wi::graphics::Shader topingBLASShader_;
    wi::graphics::Shader shadowShader_;
    wi::graphics::Shader aoBlurShader_;
    wi::graphics::Shader aoApplyShader_;

    // Blocky BLAS resources
    mutable wi::graphics::GPUBuffer blasPositionBuffer_;
    wi::graphics::GPUBuffer blasIndexBuffer_;
    mutable wi::graphics::RaytracingAccelerationStructure blockyBLAS_;
    mutable uint32_t blockyBLASCapacity_ = 0;
    mutable uint32_t blockyVertexCount_ = 0;

    // Smooth BLAS
    mutable wi::graphics::RaytracingAccelerationStructure smoothBLAS_;
    mutable uint32_t smoothBLASCapacity_ = 0;
    mutable uint32_t smoothVertexCount_ = 0;

    // Toping BLAS
    mutable wi::graphics::RaytracingAccelerationStructure topingBLAS_;
    mutable uint32_t topingBLASASCapacity_ = 0;
    mutable uint32_t topingVertexCount_ = 0;
    mutable DeferredGPUBuffer topingBLASPositionBuf_;
    mutable wi::graphics::GPUBuffer topingBLASIndexBuffer_;
    mutable uint32_t topingBLASIndexCount_ = 0;
    wi::graphics::GPUBuffer topingBLASGroupBuffer_;

    // TLAS
    mutable wi::graphics::RaytracingAccelerationStructure tlas_;
    mutable uint32_t tlasInstanceCount_ = 0;

    uint32_t maxBlasVertices_ = 0;
};

} // namespace voxel