/*
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 "../lib/stb_truetype.h"
#ifdef CRAFTER_GRAPHICS_RENDERER_VULKAN
#include <vulkan/vulkan.h>
#endif
export module Crafter.Graphics:RenderingElement2DVulkan;
#ifdef CRAFTER_GRAPHICS_RENDERER_VULKAN
import Crafter.Asset;
import std;
import :Transform2D;
import :VulkanBuffer;
import :Types;
import :Window;
import :DescriptorHeapVulkan;

export namespace Crafter {
    struct __attribute__((packed)) RenderingElement2DVulkanTransformInfo {
        ScaleData2D scaled;     // 0 - 8 bytes
        std::uint16_t bufferX;  // 8 - 2 bytes
        std::uint16_t bufferY;  // 10 - 2 bytes
        //12 bytes total;
    };

    struct RenderingElement2DVulkanBase : Transform2D {
        std::uint16_t index;
        std::uint16_t bufferX;
        std::uint16_t bufferY;
        std::array<VulkanBufferBase*, Window::numFrames> buffers;
        RenderingElement2DVulkanBase(Anchor2D anchor) : Transform2D(anchor) {

        }
        RenderingElement2DVulkanBase(Anchor2D anchor, std::uint16_t bufferX, std::uint16_t bufferY) : bufferX(bufferX), bufferY(bufferY), Transform2D(anchor) {

        }
        RenderingElement2DVulkanBase(Anchor2D anchor, std::uint16_t bufferX, std::uint16_t bufferY, std::array<VulkanBufferBase*, Window::numFrames>&& buffers) : bufferX(bufferX), bufferY(bufferY), buffers(std::move(buffers)), Transform2D(anchor) {

        }
    };

    std::vector<RenderingElement2DVulkanBase*> renderingElement2DVulkans;
    VulkanBuffer<RenderingElement2DVulkanTransformInfo, true> renderingElement2DVulkanTransformBuffer[Window::numFrames];

    template<bool Owning, bool Mapped>
    struct RenderingElement2DVulkan : RenderingElement2DVulkanBase {
        RenderingElement2DVulkan(Anchor2D anchor, RendertargetBase<Window::numFrames>& target, Transform2D& parent) requires(Owning) : RenderingElement2DVulkanBase(anchor) {
            renderingElement2DVulkans.push_back(this);
            UpdatePosition(target, parent);
            this->bufferX = this->scaled.size.x;
            this->bufferY = this->scaled.size.y;
            for(std::uint8_t i = 0; i < Window::numFrames; i++) {
                buffers[i] = new VulkanBuffer<Vector<_Float16, 4, 4>, Mapped>(); 
                reinterpret_cast<VulkanBuffer<Vector<_Float16, 4, 4>, Mapped>*>(buffers[i])->Create(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, bufferX * bufferY);
            }
        }
        RenderingElement2DVulkan(Anchor2D anchor, std::uint16_t bufferX, std::uint16_t bufferY) requires(Owning) : RenderingElement2DVulkanBase(anchor, bufferX, bufferY) {
            renderingElement2DVulkans.push_back(this);
            for(std::uint8_t i = 0; i < Window::numFrames; i++) {
                buffers[i] = new VulkanBuffer<Vector<_Float16, 4, 4>, Mapped>(); 
                reinterpret_cast<VulkanBuffer<Vector<_Float16, 4, 4>, Mapped>*>(buffers[i])->Create(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, bufferX * bufferY);
            }
        }
        RenderingElement2DVulkan(Anchor2D anchor, std::uint16_t bufferX, std::uint16_t bufferY, std::array<VulkanBufferBase*, Window::numFrames>&& buffers) requires(!Owning) : RenderingElement2DVulkanBase(anchor, bufferX, bufferY, std::move(buffers)) {
            renderingElement2DVulkans.push_back(this);
        }

        ~RenderingElement2DVulkan() {
            if constexpr(Owning) {
                for(VulkanBufferBase* buffer : buffers) {
                    delete reinterpret_cast<VulkanBuffer<Vector<_Float16, 4, 4>, Mapped>*>(buffer);
                }
            }
            auto it = std::find(renderingElement2DVulkans.begin(), renderingElement2DVulkans.end(), this);
            if (it != renderingElement2DVulkans.end()) {
                renderingElement2DVulkans.erase(it);
            }
        }
    
        RenderingElement2DVulkan(RenderingElement2DVulkan&) = delete;
        RenderingElement2DVulkan& operator=(RenderingElement2DVulkan&) = delete;

        void UpdatePosition(RendertargetBase<Window::numFrames>& window, Transform2D& parent) override {
            this->ScaleElement(parent);
            RenderingElement2DVulkanTransformInfo* val = reinterpret_cast<RenderingElement2DVulkanTransformInfo*>(reinterpret_cast<char*>(renderingElement2DVulkanTransformBuffer[window.frame].value) + sizeof(RenderingElement2DVulkanTransformInfo));
            val[index].scaled =  this->scaled;
            for(Transform2D* child : this->children) {
                child->UpdatePosition(window, *this);
            }
        }
    };

    void WriteRenderingElement2DVulkanDescriptors(std::span<VkResourceDescriptorInfoEXT> infos, std::span<VkHostAddressRangeEXT> ranges, std::uint16_t start, std::uint32_t bufferOffset, DescriptorHeapVulkan& descriptorHeap) {
        VkDeviceAddressRangeKHR transformRanges[Window::numFrames] = {
            {
                .address = renderingElement2DVulkanTransformBuffer[0].address,
                .size = renderingElement2DVulkanTransformBuffer[0].size
            },
            {
                .address = renderingElement2DVulkanTransformBuffer[1].address,
                .size = renderingElement2DVulkanTransformBuffer[1].size
            },
            {
                .address = renderingElement2DVulkanTransformBuffer[2].address,
                .size = renderingElement2DVulkanTransformBuffer[2].size
            }
        };
        
        for(std::uint8_t i = 0; i < Window::numFrames; i++) {
            ranges[start + i] = {
                .address = descriptorHeap.resourceHeap[i].value + bufferOffset,
                .size = Device::descriptorHeapProperties.bufferDescriptorSize
            };
            infos[start + i] = {
                .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
                .data = { .pAddressRange =  &transformRanges[i]}
            };
        }

        start += 3;
        bufferOffset += Device::descriptorHeapProperties.bufferDescriptorSize;

        std::vector<VkDeviceAddressRangeKHR> bufferRanges(renderingElement2DVulkans.size() * Window::numFrames);

        std::uint16_t rangeOffset = 0;

        for(std::uint8_t i2 = 0; i2 < Window::numFrames; i2++) {
            for(std::uint16_t i = 0; i < renderingElement2DVulkans.size(); i++) {
                ranges[start + i] = {
                    .address = descriptorHeap.resourceHeap[i2].value + bufferOffset + Device::descriptorHeapProperties.bufferDescriptorSize * i,
                    .size = Device::descriptorHeapProperties.bufferDescriptorSize
                };
                infos[start + i] = {
                    .sType = VK_STRUCTURE_TYPE_RESOURCE_DESCRIPTOR_INFO_EXT,
                    .type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
                    .data = { .pAddressRange =  &bufferRanges[i]}
                };
                bufferRanges[rangeOffset + i] =  {
                    .address = renderingElement2DVulkans[i]->buffers[i2]->address,
                    .size = renderingElement2DVulkans[i]->buffers[i2]->size
                };
            }
            start += renderingElement2DVulkans.size();
            rangeOffset += renderingElement2DVulkans.size();
        }

        Device::vkWriteResourceDescriptorsEXT(Device::device, start, infos.data(), ranges.data());

        for(std::uint8_t i = 0; i < Window::numFrames; i++) {
            descriptorHeap.resourceHeap[i].FlushDevice();
        }
    }

    void InitializeRenderingElement2DVulkanBuffer() {
        for(std::uint8_t frame = 0; frame < Window::numFrames; frame++) {
            renderingElement2DVulkanTransformBuffer[frame].Create(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, renderingElement2DVulkans.size()+1);
            RenderingElement2DVulkanTransformInfo* val =  reinterpret_cast<RenderingElement2DVulkanTransformInfo*>(reinterpret_cast<char*>(renderingElement2DVulkanTransformBuffer[frame].value) + sizeof(RenderingElement2DVulkanTransformInfo));
            std::uint16_t* sizePtr =  reinterpret_cast<std::uint16_t*>(renderingElement2DVulkanTransformBuffer[frame].value);
            *sizePtr = renderingElement2DVulkans.size();
            for(std::uint16_t i = 0; i < renderingElement2DVulkans.size(); i++) {
                val[i].scaled = renderingElement2DVulkans[i]->scaled;
                val[i].bufferX = renderingElement2DVulkans[i]->bufferX;
                val[i].bufferY = renderingElement2DVulkans[i]->bufferY;
                renderingElement2DVulkans[i]->index = i;
            }
        }
    }

    void UpdateRenderingElement2DVulkanBuffer(std::uint8_t frame) {
        renderingElement2DVulkanTransformBuffer[frame].Resize(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_2_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, renderingElement2DVulkans.size()+1);
        RenderingElement2DVulkanTransformInfo* val =  reinterpret_cast<RenderingElement2DVulkanTransformInfo*>(reinterpret_cast<char*>(renderingElement2DVulkanTransformBuffer[frame].value) + sizeof(RenderingElement2DVulkanTransformInfo));
        std::uint16_t* sizePtr =  reinterpret_cast<std::uint16_t*>(renderingElement2DVulkanTransformBuffer[frame].value);
        *sizePtr = renderingElement2DVulkans.size();
        for(std::uint16_t i = 0; i < renderingElement2DVulkans.size(); i++) {
            val[i].scaled = renderingElement2DVulkans[i]->scaled;
            val[i].bufferX = renderingElement2DVulkans[i]->bufferX;
            val[i].bufferY = renderingElement2DVulkans[i]->bufferY;
            renderingElement2DVulkans[i]->index = i + 1;
        }
    }
}
#endif