package net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.VkModel;

import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandPool;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.GPUSynchronisation.TransferBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.Queues.Queue;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.system.MemoryUtil;
import org.tinylog.Logger;

import java.io.BufferedInputStream;
import java.io.DataInputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.nio.FloatBuffer;
import java.nio.IntBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.lwjgl.vulkan.VK13.*;

public class ModelsCache {
    private final Map<String, VulkanModel> ModelsMap;
    public ModelsCache(){
        ModelsMap = new HashMap<>();
    }

    public void loadModels(VulkanContext VkCtx, List<ModelData> Models, CommandPool commandPool, Queue queue){
        try {
            List<VulkanBuffer> StagingBufferList = new ArrayList<>();
            var Command = new CommandBuffer(VkCtx, commandPool, true, true);
            Command.BeginRecording();

            for (ModelData modelData : Models) {
                VulkanModel VKModel = new VulkanModel(modelData.ID());
                ModelsMap.put(VKModel.GetID(), VKModel);

                DataInputStream VertexInput = new DataInputStream(new BufferedInputStream(
                        new FileInputStream(modelData.VertexPath())));
                DataInputStream IndicesInput = new DataInputStream(new BufferedInputStream(
                        new FileInputStream(modelData.IndexPath())));

                for (MeshData meshData : modelData.Meshes()) {
                    TransferBuffer VerticesBuffers = CreateVerticesBuffer(VkCtx, meshData,VertexInput);
                    TransferBuffer IndicesBuffers = CreateIndicesBuffer(VkCtx, meshData,IndicesInput);
                    StagingBufferList.add(VerticesBuffers.SrcBuffer());
                    StagingBufferList.add(IndicesBuffers.SrcBuffer());
                    VerticesBuffers.RecordTransferCommand(Command);
                    IndicesBuffers.RecordTransferCommand(Command);
                    Logger.debug("Creating new Vulkan Mesh -> ID=[{}] MaterialID=[{}]",meshData.ID(),meshData.MaterialID());
                    VulkanMesh VkMesh = new VulkanMesh(meshData.ID(), VerticesBuffers.DstBuffer(), IndicesBuffers.DstBuffer(),
                            meshData.IndexSize() / VulkanUtils.INT_SIZE, meshData.MaterialID());
                    VKModel.GetVkMeshList().add(VkMesh);
                }
            }
        Command.EndRecording();
        Command.SubmitAndWait(VkCtx,queue);
        Command.cleanup(VkCtx,commandPool);

        StagingBufferList.forEach(b -> b.cleanup(VkCtx));
        } catch (Exception exception){
            throw new RuntimeException(exception);
        }
    }

    private static TransferBuffer CreateVerticesBuffer(VulkanContext VkCtx, MeshData meshData, DataInputStream VertexStream) throws IOException {
        Logger.debug("Creating Vertices Transfer Buffer, Mesh Data -> ID=[{}] VertexSize=[{}]",meshData.ID(),meshData.VertexSize());
        int BufferSize = meshData.VertexSize();
        var SrcBuffer = new VulkanBuffer(VkCtx,BufferSize,
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
        var DstBuffer = new VulkanBuffer(VkCtx, BufferSize,
                VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        if (BufferSize <= 0) {
            Logger.error("Computed buffer allocation size cannot be 0. SrcBuffer -> [{}] DstBuffer -> [{}]",SrcBuffer,DstBuffer);
            throw new RuntimeException("Computed buffer allocation size cannot be 0. SrcBuffer -> ["+SrcBuffer+"] DstBuffer -> ["+DstBuffer+"]");
        }
        long MappedMemory = SrcBuffer.MapMemory(VkCtx);
        FloatBuffer data = MemoryUtil.memFloatBuffer(MappedMemory, (int) SrcBuffer.GetRequestedSize());

        int ValuesToRead = meshData.VertexSize()/VulkanUtils.FLOAT_SIZE;
        if (ValuesToRead <= 0) {
            Logger.error("Cannot allocate vertex buffer: Model mesh has 0 or negative vertices! Vertex Stream -> [{}]",VertexStream);
            throw new RuntimeException("Cannot allocate vertex buffer: Model mesh has 0 or negative vertices! Vertex Stream -> [{"+VertexStream+"]");
        }
        while(ValuesToRead > 0){
            data.put(VertexStream.readFloat());
            ValuesToRead--;
        }
        SrcBuffer.UnMapMemory(VkCtx);
        return new TransferBuffer(SrcBuffer,DstBuffer);
    }

    private static TransferBuffer CreateIndicesBuffer(VulkanContext VkCtx, MeshData meshData, DataInputStream IndexStream) throws IOException{
        int BufferSize = meshData.IndexSize();
        var SrcBuffer = new VulkanBuffer(VkCtx,BufferSize,
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
        var DstBuffer = new VulkanBuffer(VkCtx, BufferSize,
                VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        long MappedMemory = SrcBuffer.MapMemory(VkCtx);
        IntBuffer data = MemoryUtil.memIntBuffer(MappedMemory,(int) SrcBuffer.GetRequestedSize());

        int ValuesToRead = meshData.IndexSize()/VulkanUtils.INT_SIZE;
        while(ValuesToRead > 0){
            data.put(IndexStream.readInt());
            ValuesToRead--;
        }
        SrcBuffer.UnMapMemory(VkCtx);
        return new TransferBuffer(SrcBuffer,DstBuffer);
    }
    public VulkanModel GetModel(String modelName){
        return ModelsMap.get(modelName);
    }
    public Map<String, VulkanModel> GetModelMap(){
        return ModelsMap;
    }
    public void CleanUp(VulkanContext VkCtx){
        ModelsMap.forEach((key,Value)->Value.CleanUp(VkCtx));
        ModelsMap.clear();
    }
}
