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

        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;

            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 (imageNext + count > imageCapacity) {
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of image slots ({} requested, {} remaining of {})", count, imageCapacity - imageNext, imageCapacity));
            }
            ImageSlotRange r{imageNext, count};
            imageNext += count;
            return r;
        }

        BufferSlotRange AllocateBufferSlots(std::uint16_t count) {
            if (bufferNext + count > bufferCapacity) {
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of buffer slots ({} requested, {} remaining of {})", count, bufferCapacity - bufferNext, bufferCapacity));
            }
            BufferSlotRange r{bufferNext, count};
            bufferNext += count;
            return r;
        }

        SamplerSlotRange AllocateSamplerSlots(std::uint16_t count) {
            if (samplerNext + count > samplerCapacity) {
                throw std::runtime_error(std::format("DescriptorHeapVulkan: out of sampler slots ({} requested, {} remaining of {})", count, samplerCapacity - samplerNext, samplerCapacity));
            }
            SamplerSlotRange r{samplerNext, count};
            samplerNext += count;
            return r;
        }

        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;
        }
    };
}