#include "VoxelRTManager.h"
#include <cstring>

using namespace wi::graphics;

namespace voxel {

void VoxelRTManager::initialize(GraphicsDevice* dev, uint32_t maxBlasVertices) {
    device_ = dev;
    maxBlasVertices_ = maxBlasVertices;

    available_ = dev->CheckCapability(GraphicsDeviceCapability::RAYTRACING);
    if (!available_) {
        wi::backlog::post("VoxelRTManager: RT not available (GPU does not support ray tracing)");
        return;
    }

    wi::renderer::LoadShader(ShaderStage::CS, blasExtractShader_, "voxel/voxelBLASExtractCS.cso");
    if (blasExtractShader_.IsValid()) {
        // BLAS position buffer: 6 float3 per quad (non-indexed triangles), raw buffer
        GPUBufferDesc posDesc;
        posDesc.size = (uint64_t)maxBlasVertices * sizeof(float) * 3;
        posDesc.bind_flags = BindFlag::UNORDERED_ACCESS | BindFlag::SHADER_RESOURCE;
        posDesc.misc_flags = ResourceMiscFlag::BUFFER_RAW;
        posDesc.stride = 0;
        posDesc.usage = Usage::DEFAULT;
        bool ok = dev->CreateBuffer(&posDesc, nullptr, &blasPositionBuffer_);

        // Sequential index buffer for BLAS
        GPUBufferDesc idxDesc;
        idxDesc.size = (uint64_t)maxBlasVertices * sizeof(uint32_t);
        idxDesc.bind_flags = BindFlag::SHADER_RESOURCE;
        idxDesc.usage = Usage::DEFAULT;
        auto fillIndices = [maxBlasVertices](void* dest) {
            uint32_t* p = (uint32_t*)dest;
            for (uint32_t i = 0; i < maxBlasVertices; i++)
                p[i] = i;
        };
        bool okIdx = dev->CreateBuffer2(&idxDesc, fillIndices, &blasIndexBuffer_);

        if (ok && blasPositionBuffer_.IsValid() && okIdx && blasIndexBuffer_.IsValid()) {
            dev->SetName(&blasPositionBuffer_, "VoxelRTManager::blasPositionBuffer");
            dev->SetName(&blasIndexBuffer_, "VoxelRTManager::blasIndexBuffer");
            wi::backlog::post("VoxelRTManager: RT available (BLAS pos "
                + std::to_string(posDesc.size / (1024*1024)) + " MB + idx "
                + std::to_string(idxDesc.size / (1024*1024)) + " MB)");
        } else {
            available_ = false;
            wi::backlog::post("VoxelRTManager: RT buffer creation failed", wi::backlog::LogLevel::Warning);
        }
    } else {
        available_ = false;
        wi::backlog::post("VoxelRTManager: BLAS extraction shader failed", wi::backlog::LogLevel::Warning);
    }

    // Toping BLAS CS
    wi::renderer::LoadShader(ShaderStage::CS, topingBLASShader_, "voxel/voxelTopingBLASCS.cso");
    if (topingBLASShader_.IsValid()) {
        static constexpr uint32_t MAX_GROUPS = 64;
        GPUBufferDesc grpDesc;
        grpDesc.size = MAX_GROUPS * 20; // 5 × uint32 per group
        grpDesc.bind_flags = BindFlag::SHADER_RESOURCE;
        grpDesc.misc_flags = ResourceMiscFlag::BUFFER_STRUCTURED;
        grpDesc.stride = 20;
        grpDesc.usage = Usage::DEFAULT;
        dev->CreateBuffer(&grpDesc, nullptr, &topingBLASGroupBuffer_);
        wi::backlog::post("VoxelRTManager: toping BLAS CS available");
    } else {
        wi::backlog::post("VoxelRTManager: toping BLAS CS failed", wi::backlog::LogLevel::Warning);
    }

    // RT Shadows + AO
    wi::renderer::LoadShader(ShaderStage::CS, shadowShader_, "voxel/voxelShadowCS.cso",
        ShaderModel::SM_6_5);
    wi::renderer::LoadShader(ShaderStage::CS, aoBlurShader_, "voxel/voxelAOBlurCS.cso");
    wi::renderer::LoadShader(ShaderStage::CS, aoApplyShader_, "voxel/voxelAOApplyCS.cso");
    if (shadowShader_.IsValid() && aoBlurShader_.IsValid() && aoApplyShader_.IsValid()) {
        shadowsEnabled_ = true;
        wi::backlog::post("VoxelRTManager: RT shadows + AO blur available");
    } else {
        wi::backlog::post("VoxelRTManager: RT shadow/AO shader(s) failed",
            wi::backlog::LogLevel::Warning);
    }
}

// ── BLAS extraction: blocky quads → float3 positions ────────────

void VoxelRTManager::dispatchBLASExtract(CommandList cmd,
    const GPUBuffer& quadBuffer,
    const GPUBuffer& chunkInfoBuffer,
    uint32_t quadCount) const
{
    if (!available_ || !blasExtractShader_.IsValid() || quadCount == 0) return;

    auto* dev = device_;

    GPUBarrier preBarriers[] = {
        GPUBarrier::Buffer(&blasPositionBuffer_,
            ResourceState::UNDEFINED, ResourceState::UNORDERED_ACCESS),
    };
    dev->Barrier(preBarriers, 1, cmd);

    dev->BindComputeShader(&blasExtractShader_, cmd);
    dev->BindResource(&quadBuffer, 0, cmd);        // t0
    dev->BindResource(&chunkInfoBuffer, 2, cmd);   // t2
    dev->BindUAV(&blasPositionBuffer_, 0, cmd);    // u0

    struct BLASPush {
        uint32_t quadCount;
        uint32_t pad[11];
    } pushData = {};
    pushData.quadCount = quadCount;
    dev->PushConstants(&pushData, sizeof(pushData), cmd);

    uint32_t groupCount = (quadCount + 63) / 64;
    dev->Dispatch(groupCount, 1, 1, cmd);

    GPUBarrier postBarriers[] = {
        GPUBarrier::Buffer(&blasPositionBuffer_,
            ResourceState::UNORDERED_ACCESS, ResourceState::SHADER_RESOURCE),
    };
    dev->Barrier(postBarriers, 1, cmd);

    blockyVertexCount_ = quadCount * 6;
}

// ── Toping BLAS extraction (GPU compute) ────────────────────────

void VoxelRTManager::dispatchTopingBLASExtract(CommandList cmd,
    const GPUBuffer& topingVertexBuffer,
    const GPUBuffer& topingInstanceBuffer,
    const void* groupsGPUData, size_t groupsGPUSize,
    uint32_t groupCount, uint32_t totalVertices) const
{
    if (!topingBLASShader_.IsValid() || !topingBLASGroupBuffer_.IsValid() ||
        !topingBLASPositionBuf_.isValid() || !topingVertexBuffer.IsValid() ||
        !topingInstanceBuffer.IsValid() || totalVertices == 0 || groupCount == 0)
        return;

    auto* dev = device_;

    // Upload group table
    dev->UpdateBuffer(&topingBLASGroupBuffer_, groupsGPUData, cmd, groupsGPUSize);

    GPUBarrier preBarriers[] = {
        GPUBarrier::Buffer(&topingBLASGroupBuffer_,
            ResourceState::COPY_DST, ResourceState::SHADER_RESOURCE),
        GPUBarrier::Buffer(&topingBLASPositionBuf_.gpu,
            ResourceState::UNDEFINED, ResourceState::UNORDERED_ACCESS),
    };
    dev->Barrier(preBarriers, 2, cmd);

    dev->BindComputeShader(&topingBLASShader_, cmd);
    dev->BindResource(&topingVertexBuffer, 4, cmd);      // t4
    dev->BindResource(&topingInstanceBuffer, 5, cmd);    // t5
    dev->BindResource(&topingBLASGroupBuffer_, 7, cmd);  // t7
    dev->BindUAV(&topingBLASPositionBuf_.gpu, 0, cmd);   // u0

    struct {
        uint32_t totalVertices;
        uint32_t groupCount;
        uint32_t pad[10];
    } pushData = {};
    pushData.totalVertices = totalVertices;
    pushData.groupCount = groupCount;
    dev->PushConstants(&pushData, sizeof(pushData), cmd);

    uint32_t threadGroups = (totalVertices + 63) / 64;
    dev->Dispatch(threadGroups, 1, 1, cmd);

    GPUBarrier postBarriers[] = {
        GPUBarrier::Buffer(&topingBLASPositionBuf_.gpu,
            ResourceState::UNORDERED_ACCESS, ResourceState::SHADER_RESOURCE),
    };
    dev->Barrier(postBarriers, 1, cmd);

    topingVertexCount_ = totalVertices;
    dirty = true;
    topingBLASDirty = false;
}

// ── Ensure toping BLAS buffer capacity ──────────────────────────

bool VoxelRTManager::ensureTopingBLASCapacity(uint32_t totalVertices) {
    if (totalVertices == 0) return false;

    bool recreated = topingBLASPositionBuf_.ensureCapacity(device_, totalVertices,
        3 * sizeof(float),
        BindFlag::UNORDERED_ACCESS | BindFlag::SHADER_RESOURCE,
        ResourceMiscFlag::BUFFER_RAW);

    if (recreated) {
        char msg[256];
        snprintf(msg, sizeof(msg), "VoxelRTManager: toping BLAS pos buffer (%u capacity, %.1f MB)",
            topingBLASPositionBuf_.capacity,
            (size_t)topingBLASPositionBuf_.capacity * 3 * sizeof(float) / (1024.0 * 1024.0));
        wi::backlog::post(msg);
    }

    // Index buffer: grow if needed
    if (topingBLASIndexCount_ < topingBLASPositionBuf_.capacity) {
        uint32_t idxCount = topingBLASPositionBuf_.capacity;
        std::vector<uint32_t> indices(idxCount);
        for (uint32_t j = 0; j < idxCount; j++) indices[j] = j;

        GPUBufferDesc idxDesc;
        idxDesc.size = (size_t)idxCount * sizeof(uint32_t);
        idxDesc.bind_flags = BindFlag::SHADER_RESOURCE;
        idxDesc.misc_flags = ResourceMiscFlag::NONE;
        idxDesc.usage = Usage::DEFAULT;
        device_->CreateBuffer(&idxDesc, indices.data(), &topingBLASIndexBuffer_);
        topingBLASIndexCount_ = idxCount;
        recreated = true;
    }

    topingBLASDirty = true;
    return recreated;
}

// ── Acceleration structure build ────────────────────────────────

void VoxelRTManager::buildAccelerationStructures(CommandList cmd,
    uint32_t buildFlags,
    const GPUBuffer& smoothVB,
    uint32_t smoothVertCount) const
{
    if (!available_) return;

    auto* dev = device_;

    // ── Blocky BLAS ──
    uint32_t blockyVertCount = blockyVertexCount_;
    if (blockyVertCount < 3) blockyVertCount = 0;
    if ((buildFlags & BUILD_BLOCKY) && blockyVertCount > 0 && blasPositionBuffer_.IsValid()) {
        if (!blockyBLAS_.IsValid() || blockyVertCount > blockyBLASCapacity_) {
            blockyBLASCapacity_ = blockyVertCount + blockyVertCount / 4;

            RaytracingAccelerationStructureDesc desc;
            desc.type = RaytracingAccelerationStructureDesc::Type::BOTTOMLEVEL;
            desc.flags = RaytracingAccelerationStructureDesc::FLAG_PREFER_FAST_BUILD;

            desc.bottom_level.geometries.resize(1);
            auto& geom = desc.bottom_level.geometries[0];
            geom.type = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::Type::TRIANGLES;
            geom.flags = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::FLAG_OPAQUE;
            geom.triangles.vertex_buffer = blasPositionBuffer_;
            geom.triangles.vertex_byte_offset = 0;
            geom.triangles.vertex_count = blockyBLASCapacity_;
            geom.triangles.vertex_stride = sizeof(float) * 3;
            geom.triangles.vertex_format = Format::R32G32B32_FLOAT;
            geom.triangles.index_buffer = blasIndexBuffer_;
            geom.triangles.index_count = blockyBLASCapacity_;
            geom.triangles.index_format = IndexBufferFormat::UINT32;
            geom.triangles.index_offset = 0;

            bool ok = dev->CreateRaytracingAccelerationStructure(&desc, &blockyBLAS_);
            if (ok) {
                dev->SetName(&blockyBLAS_, "VoxelRTManager::blockyBLAS");
                wi::backlog::post("VoxelRTManager: blocky BLAS created (capacity "
                    + std::to_string(blockyBLASCapacity_ / 3) + " tris)");
            } else {
                wi::backlog::post("VoxelRTManager: failed to create blocky BLAS", wi::backlog::LogLevel::Error);
                available_ = false;
                return;
            }
        }

        blockyBLAS_.desc.bottom_level.geometries[0].triangles.vertex_count = blockyVertCount;
        blockyBLAS_.desc.bottom_level.geometries[0].triangles.index_count = blockyVertCount;
        dev->BuildRaytracingAccelerationStructure(&blockyBLAS_, cmd, nullptr);
    }

    // ── Smooth BLAS ──
    if (smoothVertCount < 3) smoothVertCount = 0;
    if ((buildFlags & BUILD_SMOOTH) && smoothVertCount > 0 && smoothVB.IsValid()) {
        if (!smoothBLAS_.IsValid() || smoothVertCount > smoothBLASCapacity_) {
            smoothBLASCapacity_ = smoothVertCount + smoothVertCount / 4;

            RaytracingAccelerationStructureDesc desc;
            desc.type = RaytracingAccelerationStructureDesc::Type::BOTTOMLEVEL;
            desc.flags = RaytracingAccelerationStructureDesc::FLAG_PREFER_FAST_BUILD;

            desc.bottom_level.geometries.resize(1);
            auto& geom = desc.bottom_level.geometries[0];
            geom.type = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::Type::TRIANGLES;
            geom.flags = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::FLAG_OPAQUE;
            geom.triangles.vertex_buffer = smoothVB;
            geom.triangles.vertex_byte_offset = 0;
            geom.triangles.vertex_count = smoothBLASCapacity_;
            geom.triangles.vertex_stride = 32;
            geom.triangles.index_buffer = blasIndexBuffer_;
            geom.triangles.index_count = smoothBLASCapacity_;
            geom.triangles.index_format = IndexBufferFormat::UINT32;
            geom.triangles.index_offset = 0;
            geom.triangles.vertex_format = Format::R32G32B32_FLOAT;

            bool ok = dev->CreateRaytracingAccelerationStructure(&desc, &smoothBLAS_);
            if (ok) {
                dev->SetName(&smoothBLAS_, "VoxelRTManager::smoothBLAS");
                wi::backlog::post("VoxelRTManager: smooth BLAS created (capacity "
                    + std::to_string(smoothBLASCapacity_ / 3) + " tris)");
            } else {
                wi::backlog::post("VoxelRTManager: failed to create smooth BLAS", wi::backlog::LogLevel::Error);
            }
        }

        if (smoothBLAS_.IsValid()) {
            smoothBLAS_.desc.bottom_level.geometries[0].triangles.vertex_count = smoothVertCount;
            smoothBLAS_.desc.bottom_level.geometries[0].triangles.index_count = smoothVertCount;
            dev->BuildRaytracingAccelerationStructure(&smoothBLAS_, cmd, nullptr);
        }

        smoothVertexCount_ = smoothVertCount;
    }

    // ── Toping BLAS ──
    uint32_t topingVertCount = topingVertexCount_;
    if ((buildFlags & BUILD_TOPING) && topingVertCount >= 3 && topingBLASPositionBuf_.isValid()) {
        if (!topingBLAS_.IsValid() || topingVertCount > topingBLASASCapacity_) {
            topingBLASASCapacity_ = topingVertCount + topingVertCount / 4;

            RaytracingAccelerationStructureDesc desc;
            desc.type = RaytracingAccelerationStructureDesc::Type::BOTTOMLEVEL;
            desc.flags = RaytracingAccelerationStructureDesc::FLAG_PREFER_FAST_BUILD;

            desc.bottom_level.geometries.resize(1);
            auto& geom = desc.bottom_level.geometries[0];
            geom.type = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::Type::TRIANGLES;
            geom.flags = RaytracingAccelerationStructureDesc::BottomLevel::Geometry::FLAG_OPAQUE;
            geom.triangles.vertex_buffer = topingBLASPositionBuf_.gpu;
            geom.triangles.vertex_byte_offset = 0;
            geom.triangles.vertex_count = topingBLASASCapacity_;
            geom.triangles.vertex_stride = sizeof(float) * 3;
            geom.triangles.vertex_format = Format::R32G32B32_FLOAT;
            geom.triangles.index_buffer = topingBLASIndexBuffer_;
            geom.triangles.index_count = topingBLASASCapacity_;
            geom.triangles.index_format = IndexBufferFormat::UINT32;
            geom.triangles.index_offset = 0;

            bool ok = dev->CreateRaytracingAccelerationStructure(&desc, &topingBLAS_);
            if (ok) {
                dev->SetName(&topingBLAS_, "VoxelRTManager::topingBLAS");
                wi::backlog::post("VoxelRTManager: toping BLAS created (capacity "
                    + std::to_string(topingBLASASCapacity_ / 3) + " tris)");
            } else {
                wi::backlog::post("VoxelRTManager: failed to create toping BLAS", wi::backlog::LogLevel::Error);
            }
        }

        if (topingBLAS_.IsValid()) {
            topingBLAS_.desc.bottom_level.geometries[0].triangles.vertex_count = topingVertCount;
            topingBLAS_.desc.bottom_level.geometries[0].triangles.index_count = topingVertCount;
            dev->BuildRaytracingAccelerationStructure(&topingBLAS_, cmd, nullptr);
        }
    }

    // Memory barrier: sync BLAS builds before TLAS
    {
        GPUBarrier barriers[] = { GPUBarrier::Memory() };
        dev->Barrier(barriers, 1, cmd);
    }

    // ── TLAS ──
    uint32_t instanceCount = 0;
    if (blockyBLAS_.IsValid()) instanceCount++;
    if (smoothBLAS_.IsValid() && smoothVertCount > 0) instanceCount++;
    if (topingBLAS_.IsValid() && topingVertCount >= 3) instanceCount++;
    if (instanceCount == 0) { dirty = false; return; }

    if (!tlas_.IsValid() || instanceCount != tlasInstanceCount_) {
        const size_t instSize = dev->GetTopLevelAccelerationStructureInstanceSize();

        auto setIdentity = [](float transform[3][4]) {
            std::memset(transform, 0, sizeof(float) * 12);
            transform[0][0] = 1.0f;
            transform[1][1] = 1.0f;
            transform[2][2] = 1.0f;
        };

        const RaytracingAccelerationStructure* blockyPtr = blockyBLAS_.IsValid() ? &blockyBLAS_ : nullptr;
        const RaytracingAccelerationStructure* smoothPtr = (smoothBLAS_.IsValid() && smoothVertCount > 0) ? &smoothBLAS_ : nullptr;
        const RaytracingAccelerationStructure* topingPtr = (topingBLAS_.IsValid() && topingVertCount >= 3) ? &topingBLAS_ : nullptr;

        RaytracingAccelerationStructureDesc desc;
        desc.flags = RaytracingAccelerationStructureDesc::FLAG_PREFER_FAST_BUILD;
        desc.type = RaytracingAccelerationStructureDesc::Type::TOPLEVEL;
        desc.top_level.count = instanceCount;

        GPUBufferDesc bufdesc;
        bufdesc.misc_flags = ResourceMiscFlag::RAY_TRACING;
        bufdesc.stride = (uint32_t)instSize;
        bufdesc.size = bufdesc.stride * desc.top_level.count;

        auto initInstances = [&](void* dest) {
            uint32_t idx = 0;
            auto addInstance = [&](const RaytracingAccelerationStructure* blas, uint32_t id) {
                if (!blas) return;
                RaytracingAccelerationStructureDesc::TopLevel::Instance inst;
                setIdentity(inst.transform);
                inst.instance_id = id;  inst.instance_mask = 0xFF;
                inst.instance_contribution_to_hit_group_index = 0;  inst.flags = 0;
                inst.bottom_level = blas;
                dev->WriteTopLevelAccelerationStructureInstance(&inst, (uint8_t*)dest + idx * instSize);
                idx++;
            };
            addInstance(blockyPtr, 0);
            addInstance(smoothPtr, 1);
            addInstance(topingPtr, 2);
        };

        bool ok = dev->CreateBuffer2(&bufdesc, initInstances, &desc.top_level.instance_buffer);
        if (!ok) {
            wi::backlog::post("VoxelRTManager: failed to create TLAS instance buffer", wi::backlog::LogLevel::Error);
            dirty = false;
            return;
        }

        ok = dev->CreateRaytracingAccelerationStructure(&desc, &tlas_);
        if (!ok) {
            wi::backlog::post("VoxelRTManager: failed to create TLAS", wi::backlog::LogLevel::Error);
            dirty = false;
            return;
        }

        tlasInstanceCount_ = instanceCount;
        wi::backlog::post("VoxelRTManager: TLAS created (" + std::to_string(instanceCount) + " instances)");
    }

    dev->BuildRaytracingAccelerationStructure(&tlas_, cmd, nullptr);

    {
        GPUBarrier barriers[] = { GPUBarrier::Memory(&tlas_) };
        dev->Barrier(barriers, 1, cmd);
    }

    dirty = false;
}

// ── RT Shadow + AO dispatch ─────────────────────────────────────

void VoxelRTManager::dispatchShadows(CommandList cmd,
    const Texture& depthBuffer,
    const Texture& renderTarget,
    const Texture& normalTarget,
    const GPUBuffer& constantBuffer) const
{
    if (!shadowsEnabled_ || !shadowShader_.IsValid() || !tlas_.IsValid())
        return;

    auto* dev = device_;
    uint32_t w = renderTarget.GetDesc().width;
    uint32_t h = renderTarget.GetDesc().height;
    uint32_t gx = (w + 7) / 8;
    uint32_t gy = (h + 7) / 8;

    // Pass 1: Shadow + raw AO
    {
        GPUBarrier preBarriers[] = {
            GPUBarrier::Image(&const_cast<Texture&>(depthBuffer),
                ResourceState::DEPTHSTENCIL, ResourceState::SHADER_RESOURCE),
            GPUBarrier::Image(&const_cast<Texture&>(renderTarget),
                ResourceState::SHADER_RESOURCE, ResourceState::UNORDERED_ACCESS),
            GPUBarrier::Image(&aoRawTexture,
                ResourceState::SHADER_RESOURCE, ResourceState::UNORDERED_ACCESS),
        };
        dev->Barrier(preBarriers, 3, cmd);

        dev->BindComputeShader(&shadowShader_, cmd);
        dev->BindResource(&depthBuffer, 0, cmd);
        dev->BindResource(&normalTarget, 1, cmd);
        dev->BindResource(&tlas_, 2, cmd);
        dev->BindResource(&aoHistoryTexture, 3, cmd);
        dev->BindUAV(&renderTarget, 0, cmd);
        dev->BindUAV(&aoRawTexture, 1, cmd);
        dev->BindConstantBuffer(&constantBuffer, 0, cmd);

        struct ShadowPush {
            uint32_t width, height;
            float normalBias, shadowMaxDist;
            uint32_t debugMode;
            float aoRadius;
            uint32_t aoRayCount;
            float aoStrength;
            uint32_t frameIndex;
            uint32_t historyValid;
            uint32_t pad[2];
        } pushData = {};
        pushData.width = w;
        pushData.height = h;
        pushData.normalBias = 0.15f;
        pushData.shadowMaxDist = 512.0f;
        pushData.debugMode = shadowDebug_;
        pushData.aoRadius = 8.0f;
        pushData.aoRayCount = 4;
        pushData.aoStrength = 0.7f;
        pushData.frameIndex = frameCounter++;
        pushData.historyValid = aoHistoryValid ? 1u : 0u;
        dev->PushConstants(&pushData, sizeof(pushData), cmd);
        dev->Dispatch(gx, gy, 1, cmd);
    }

    // Pass 1.5: Copy raw AO → history
    {
        GPUBarrier copyBarriers[] = {
            GPUBarrier::Image(&aoRawTexture,
                ResourceState::UNORDERED_ACCESS, ResourceState::COPY_SRC),
            GPUBarrier::Image(&aoHistoryTexture,
                ResourceState::SHADER_RESOURCE, ResourceState::COPY_DST),
        };
        dev->Barrier(copyBarriers, 2, cmd);
        dev->CopyResource(&aoHistoryTexture, &aoRawTexture, cmd);

        GPUBarrier postCopyBarriers[] = {
            GPUBarrier::Image(&aoRawTexture,
                ResourceState::COPY_SRC, ResourceState::SHADER_RESOURCE),
            GPUBarrier::Image(&aoHistoryTexture,
                ResourceState::COPY_DST, ResourceState::SHADER_RESOURCE),
        };
        dev->Barrier(postCopyBarriers, 2, cmd);
        aoHistoryValid = true;
    }

    // Pass 2: Bilateral blur horizontal
    {
        GPUBarrier barriers[] = {
            GPUBarrier::Image(&aoBlurredTexture,
                ResourceState::SHADER_RESOURCE, ResourceState::UNORDERED_ACCESS),
        };
        dev->Barrier(barriers, 1, cmd);

        dev->BindComputeShader(&aoBlurShader_, cmd);
        dev->BindResource(&aoRawTexture, 0, cmd);
        dev->BindResource(&depthBuffer, 1, cmd);
        dev->BindResource(&normalTarget, 2, cmd);
        dev->BindUAV(&aoBlurredTexture, 0, cmd);

        struct BlurPush {
            uint32_t width, height, direction, radius;
            float depthThreshold, normalThreshold;
            uint32_t pad[6];
        } blurPush = {};
        blurPush.width = w; blurPush.height = h;
        blurPush.direction = 0; blurPush.radius = 6;
        blurPush.depthThreshold = 0.001f; blurPush.normalThreshold = 0.9f;
        dev->PushConstants(&blurPush, sizeof(blurPush), cmd);
        dev->Dispatch(gx, gy, 1, cmd);
    }

    // Pass 3: Bilateral blur vertical
    {
        GPUBarrier barriers[] = {
            GPUBarrier::Image(&aoBlurredTexture,
                ResourceState::UNORDERED_ACCESS, ResourceState::SHADER_RESOURCE),
            GPUBarrier::Image(&aoRawTexture,
                ResourceState::SHADER_RESOURCE, ResourceState::UNORDERED_ACCESS),
        };
        dev->Barrier(barriers, 2, cmd);

        dev->BindComputeShader(&aoBlurShader_, cmd);
        dev->BindResource(&aoBlurredTexture, 0, cmd);
        dev->BindResource(&depthBuffer, 1, cmd);
        dev->BindResource(&normalTarget, 2, cmd);
        dev->BindUAV(&aoRawTexture, 0, cmd);

        struct BlurPush {
            uint32_t width, height, direction, radius;
            float depthThreshold, normalThreshold;
            uint32_t pad[6];
        } blurPush = {};
        blurPush.width = w; blurPush.height = h;
        blurPush.direction = 1; blurPush.radius = 6;
        blurPush.depthThreshold = 0.001f; blurPush.normalThreshold = 0.9f;
        dev->PushConstants(&blurPush, sizeof(blurPush), cmd);
        dev->Dispatch(gx, gy, 1, cmd);
    }

    // Pass 4: Apply blurred AO
    {
        GPUBarrier barriers[] = {
            GPUBarrier::Image(&aoRawTexture,
                ResourceState::UNORDERED_ACCESS, ResourceState::SHADER_RESOURCE),
        };
        dev->Barrier(barriers, 1, cmd);

        dev->BindComputeShader(&aoApplyShader_, cmd);
        dev->BindResource(&aoRawTexture, 0, cmd);
        dev->BindResource(&depthBuffer, 1, cmd);
        dev->BindUAV(&renderTarget, 0, cmd);

        struct ApplyPush {
            uint32_t width, height, debugMode;
            uint32_t pad[9];
        } applyPush = {};
        applyPush.width = w; applyPush.height = h;
        applyPush.debugMode = shadowDebug_;
        dev->PushConstants(&applyPush, sizeof(applyPush), cmd);
        dev->Dispatch(gx, gy, 1, cmd);
    }

    // Restore resource states
    GPUBarrier postBarriers[] = {
        GPUBarrier::Image(&const_cast<Texture&>(depthBuffer),
            ResourceState::SHADER_RESOURCE, ResourceState::DEPTHSTENCIL),
        GPUBarrier::Image(&const_cast<Texture&>(renderTarget),
            ResourceState::UNORDERED_ACCESS, ResourceState::SHADER_RESOURCE),
    };
    dev->Barrier(postBarriers, 2, cmd);
}

} // namespace voxel
