/*
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
*/

module;
#include "vulkan/vulkan.h"
export module Crafter.Graphics:DescriptorHeapVulkan;
import std;
import :Device;
import :Window;
import :Types;
import :VulkanBuffer;

export namespace Crafter {
    struct ImageSlotRange   { std::uint16_t firstElement; std::uint16_t count; };
    struct BufferSlotRange  { std::uint16_t firstElement; std::uint16_t count; };
    struct SamplerSlotRange { std::uint16_t firstElement; std::uint16_t count; };

    struct DescriptorHeapVulkan {
        VulkanBuffer<std::uint8_t, true> resourceHeap[Window::numFrames];
        VulkanBuffer<std::uint8_t, true> samplerHeap[Window::numFrames];
        std::uint32_t bufferStartOffset;
        std::uint16_t bufferStartElement;

        std::uint16_t imageCapacity = 0;
        std::uint16_t bufferCapacity = 0;
        std::uint16_t samplerCapacity = 0;
        std::uint16_t imageNext = 0;
        std::uint16_t bufferNext = 0;
        std::uint16_t samplerNext = 0;

        // Per-pool freelists of recyclable slot indices (raw, pre-bufferStartElement
        // for the buffer pool). Populated by Free*Slots, consumed by Allocate*Slots
        // before falling through to the bump pointer. Single-slot allocations only —
        // multi-slot allocations need contiguity and bypass the freelist.
        std::vector<std::uint16_t> imageFreelist;
        std::vector<std::uint16_t> bufferFreelist;
        std::vector<std::uint16_t> samplerFreelist;

        void Initialize(std::uint16_t images, std::uint16_t buffers, std::uint16_t samplers) {
            std::uint32_t descriptorRegion = images * Device::descriptorHeapProperties.imageDescriptorSize + buffers * Device::descriptorHeapProperties.bufferDescriptorSize;
            std::uint32_t alignedDescriptorRegion = (descriptorRegion + Device::descriptorHeapProperties.imageDescriptorAlignment - 1) & ~(Device::descriptorHeapProperties.imageDescriptorAlignment - 1);
            std::uint32_t resourceSize = alignedDescriptorRegion + Device::descriptorHeapProperties.minResourceHeapReservedRange;
            std::uint32_t samplerSize = samplers * Device::descriptorHeapProperties.samplerDescriptorSize + Device::descriptorHeapProperties.minSamplerHeapReservedRange;
            bufferStartElement = images * Device::descriptorHeapProperties.imageDescriptorSize / Device::descriptorHeapProperties.bufferDescriptorSize;

            if(images > 0 && bufferStartElement == 0) {
                bufferStartElement = 1;
            }
            bufferStartOffset = bufferStartElement * Device::descriptorHeapProperties.bufferDescriptorSize;

            imageCapacity = images;
            bufferCapacity = buffers;
            samplerCapacity = samplers;
            imageNext = 0;
            bufferNext = 0;
            samplerNext = 0;
            // Reserve so Free*Slots' push_back is noexcept (slot handle dtors
            // call Free*Slots and must not throw). At most `capacity` slots
            // can ever be free at one time.
            imageFreelist.clear();
            imageFreelist.reserve(images);
            bufferFreelist.clear();
            bufferFreelist.reserve(buffers);
            samplerFreelist.clear();
            samplerFreelist.reserve(samplers);

            for(std::uint8_t i = 0; i < Window::numFrames; i++) {
                resourceHeap[i].Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_DESCRIPTOR_HEAP_BIT_EXT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, resourceSize);
                samplerHeap[i].Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_DESCRIPTOR_HEAP_BIT_EXT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, samplerSize);
            }
        }

        ImageSlotRange AllocateImageSlots(std::uint16_t count) {
            if (count == 1 && !imageFreelist.empty()) {
                std::uint16_t s = imageFreelist.back();
                imageFreelist.pop_back();
                return {s, 1};
            }
            if (imageNext + count > imageCapacity) {
                std::uint16_t remaining = static_cast<std::uint16_t>((imageCapacity - imageNext) + imageFreelist.size());
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of image slots ({} requested, {} remaining of {})", count, remaining, imageCapacity));
            }
            ImageSlotRange r{imageNext, count};
            imageNext += count;
            return r;
        }

        BufferSlotRange AllocateBufferSlots(std::uint16_t count) {
            if (count == 1 && !bufferFreelist.empty()) {
                std::uint16_t s = bufferFreelist.back();
                bufferFreelist.pop_back();
                return {s, 1};
            }
            if (bufferNext + count > bufferCapacity) {
                std::uint16_t remaining = static_cast<std::uint16_t>((bufferCapacity - bufferNext) + bufferFreelist.size());
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of buffer slots ({} requested, {} remaining of {})", count, remaining, bufferCapacity));
            }
            BufferSlotRange r{bufferNext, count};
            bufferNext += count;
            return r;
        }

        SamplerSlotRange AllocateSamplerSlots(std::uint16_t count) {
            if (count == 1 && !samplerFreelist.empty()) {
                std::uint16_t s = samplerFreelist.back();
                samplerFreelist.pop_back();
                return {s, 1};
            }
            if (samplerNext + count > samplerCapacity) {
                std::uint16_t remaining = static_cast<std::uint16_t>((samplerCapacity - samplerNext) + samplerFreelist.size());
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of sampler slots ({} requested, {} remaining of {})", count, remaining, samplerCapacity));
            }
            SamplerSlotRange r{samplerNext, count};
            samplerNext += count;
            return r;
        }

        // Return slots to the per-pool freelist. The descriptors at these slots
        // are NOT zeroed; the next allocation that reuses a slot overwrites it
        // via WriteSampledImageDescriptor / WriteBufferDescriptor / sampler write.
        // Caller MUST ensure no in-flight GPU frame still references the freed
        // slots — typically vkDeviceWaitIdle (or a scene-switch barrier) before
        // freeing. Multi-slot ranges are decomposed into their individual slots.
        // noexcept: capacity is reserved at Initialize so push_back never reallocates.
        void FreeImageSlots(ImageSlotRange r) noexcept {
            for (std::uint16_t i = 0; i < r.count; ++i) {
                imageFreelist.push_back(static_cast<std::uint16_t>(r.firstElement + i));
            }
        }

        void FreeBufferSlots(BufferSlotRange r) noexcept {
            for (std::uint16_t i = 0; i < r.count; ++i) {
                bufferFreelist.push_back(static_cast<std::uint16_t>(r.firstElement + i));
            }
        }

        void FreeSamplerSlots(SamplerSlotRange r) noexcept {
            for (std::uint16_t i = 0; i < r.count; ++i) {
                samplerFreelist.push_back(static_cast<std::uint16_t>(r.firstElement + i));
            }
        }

        std::uint32_t ImageByteOffset(std::uint16_t firstElement) const {
            return firstElement * Device::descriptorHeapProperties.imageDescriptorSize;
        }
        std::uint32_t BufferByteOffset(std::uint16_t firstElement) const {
            return bufferStartOffset + firstElement * Device::descriptorHeapProperties.bufferDescriptorSize;
        }
        std::uint32_t SamplerByteOffset(std::uint16_t firstElement) const {
            return firstElement * Device::descriptorHeapProperties.samplerDescriptorSize;
        }

        inline static std::uint32_t GetBufferOffset(std::uint16_t images, std::uint16_t buffers) {
            std::uint32_t bufferStartElement = images * Device::descriptorHeapProperties.imageDescriptorSize / Device::descriptorHeapProperties.bufferDescriptorSize;

            if(images > 0 && bufferStartElement == 0) {
                bufferStartElement = 1;
            }
            return bufferStartElement * Device::descriptorHeapProperties.bufferDescriptorSize;
        }
        inline static std::uint16_t GetBufferOffsetElement(std::uint16_t images, std::uint16_t buffers) {
            std::uint16_t bufferStartElement = images * Device::descriptorHeapProperties.imageDescriptorSize / Device::descriptorHeapProperties.bufferDescriptorSize;

            if(images > 0 && bufferStartElement == 0) {
                bufferStartElement = 1;
            }

            return bufferStartElement;
        }
    };

    // ─── RAII slot handles ──────────────────────────────────────────────
    //
    // Move-only owners of a single bindless slot. Destructor returns the slot
    // to the heap's freelist (noexcept — Free*Slots can't allocate because
    // Initialize reserved the freelist). Implicitly convertible to std::uint16_t
    // for use with FillHeader / WriteFooDescriptor / etc.; for buffer slots
    // this conversion folds in the heap's bufferStartElement so callers always
    // get the absolute heap index that GLSL `descriptor_heap` expects.
    //
    // Lifetime: the handle MUST be destroyed (or moved-from then destroyed)
    // before the DescriptorHeapVulkan it references. The handle MUST NOT be
    // destroyed while a frame in-flight on the GPU still references its slot —
    // free at scene-switch boundaries (vkDeviceWaitIdle), not mid-frame.
    //
    // An empty (default-constructed or moved-from) handle is a no-op on
    // destruction and converts to the sentinel 0xFFFF.

    struct ImageSlot {
        ImageSlot() = default;
        ImageSlot(DescriptorHeapVulkan* heap, std::uint16_t raw) noexcept
            : heap_(heap), raw_(raw) {}

        ImageSlot(const ImageSlot&) = delete;
        ImageSlot& operator=(const ImageSlot&) = delete;

        ImageSlot(ImageSlot&& o) noexcept : heap_(o.heap_), raw_(o.raw_) {
            o.heap_ = nullptr;
        }
        ImageSlot& operator=(ImageSlot&& o) noexcept {
            if (this != &o) {
                Reset();
                heap_ = o.heap_;
                raw_  = o.raw_;
                o.heap_ = nullptr;
            }
            return *this;
        }

        ~ImageSlot() noexcept { Reset(); }

        void Reset() noexcept {
            if (heap_) {
                heap_->FreeImageSlots({raw_, 1});
                heap_ = nullptr;
            }
        }

        operator std::uint16_t() const noexcept {
            return heap_ ? raw_ : std::uint16_t{0xFFFF};
        }
        explicit operator bool() const noexcept { return heap_ != nullptr; }

    private:
        DescriptorHeapVulkan* heap_ = nullptr;
        std::uint16_t         raw_  = 0;
    };

    struct BufferSlot {
        BufferSlot() = default;
        BufferSlot(DescriptorHeapVulkan* heap, std::uint16_t raw) noexcept
            : heap_(heap), raw_(raw) {}

        BufferSlot(const BufferSlot&) = delete;
        BufferSlot& operator=(const BufferSlot&) = delete;

        BufferSlot(BufferSlot&& o) noexcept : heap_(o.heap_), raw_(o.raw_) {
            o.heap_ = nullptr;
        }
        BufferSlot& operator=(BufferSlot&& o) noexcept {
            if (this != &o) {
                Reset();
                heap_ = o.heap_;
                raw_  = o.raw_;
                o.heap_ = nullptr;
            }
            return *this;
        }

        ~BufferSlot() noexcept { Reset(); }

        void Reset() noexcept {
            if (heap_) {
                heap_->FreeBufferSlots({raw_, 1});
                heap_ = nullptr;
            }
        }

        // Absolute heap index = bufferStartElement + raw. GLSL descriptor_heap
        // SSBOs are indexed in buffer-descriptor units from heap byte 0; the
        // raw bump-allocator index is relative to the buffer region's start.
        operator std::uint16_t() const noexcept {
            return heap_
                ? static_cast<std::uint16_t>(heap_->bufferStartElement + raw_)
                : std::uint16_t{0xFFFF};
        }
        explicit operator bool() const noexcept { return heap_ != nullptr; }

    private:
        DescriptorHeapVulkan* heap_ = nullptr;
        std::uint16_t         raw_  = 0;
    };

    struct SamplerSlot {
        SamplerSlot() = default;
        SamplerSlot(DescriptorHeapVulkan* heap, std::uint16_t raw) noexcept
            : heap_(heap), raw_(raw) {}

        SamplerSlot(const SamplerSlot&) = delete;
        SamplerSlot& operator=(const SamplerSlot&) = delete;

        SamplerSlot(SamplerSlot&& o) noexcept : heap_(o.heap_), raw_(o.raw_) {
            o.heap_ = nullptr;
        }
        SamplerSlot& operator=(SamplerSlot&& o) noexcept {
            if (this != &o) {
                Reset();
                heap_ = o.heap_;
                raw_  = o.raw_;
                o.heap_ = nullptr;
            }
            return *this;
        }

        ~SamplerSlot() noexcept { Reset(); }

        void Reset() noexcept {
            if (heap_) {
                heap_->FreeSamplerSlots({raw_, 1});
                heap_ = nullptr;
            }
        }

        operator std::uint16_t() const noexcept {
            return heap_ ? raw_ : std::uint16_t{0xFFFF};
        }
        explicit operator bool() const noexcept { return heap_ != nullptr; }

    private:
        DescriptorHeapVulkan* heap_ = nullptr;
        std::uint16_t         raw_  = 0;
    };
}