package net.halbear.Terrain4J.EngineCore.Vulkan.Rendering.Pipeline.Images;

import net.halbear.Terrain4J.EngineCore.Logic.Rendering.VulkanContext;
import net.halbear.Terrain4J.EngineCore.Main.Util.Maths.T4Math;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DeviceLayers.CommandBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Structure.DisplayToScreen.ImageView;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanBuffer;
import net.halbear.Terrain4J.EngineCore.Vulkan.Util.VulkanUtils;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.*;

import java.nio.ByteBuffer;

import static org.lwjgl.vulkan.VK10.*;
import static org.lwjgl.vulkan.VK13.VK_ACCESS_2_NONE;
import static org.lwjgl.vulkan.VK13.vkCmdPipelineBarrier2;

public class Texture {

    private final String ID;
    private final int Height;
    private final int Width;
    private final Image  image;
    private final ImageView imageView;
    private boolean RecordedTransition;
    private boolean Transparent;
    private VulkanBuffer stgBuffer;

    public Texture(VulkanContext VkCtx, String ID, ImageSrc imageSrc, int ImageFormat){
        this.ID = ID;
        RecordedTransition = false;
        Width = imageSrc.Width();
        Height = imageSrc.Height();
        SetTransparency(imageSrc.data());
        int MipMapLevels = (int) Math.floor(T4Math.Log2(Math.min(Width, Height))) + 1;
        CreateStgBuffer(VkCtx, imageSrc.data());
        var ImageData = new Image.ImageData().Width(Width).Height(Height)
                .Usage(VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
                .Format(ImageFormat).MipMapLevels(MipMapLevels);
        image = new Image(VkCtx, ImageData);
        var ImageViewData = new ImageView.ImageViewData().Format(image.GetFormat())
                .AspectMask(VK_IMAGE_ASPECT_COLOR_BIT).MipLevels(MipMapLevels);
        imageView = new ImageView(VkCtx.GetDevice(), image.getVulkanImage(), ImageViewData, false);
    }
    public boolean HasTransparency(){return Transparent;}
    public void SetTransparency(ByteBuffer data){
        int PixelCount = data.capacity()/4;
        int Offset = 0;
        Transparent = false;
        for(int i = 0; i < PixelCount; i++){
            int alpha = (0xFF & data.get(Offset + 3));
            if( alpha < 255){
                Transparent = true;
            }
            Offset += 4;
        }
    }
    public void CreateStgBuffer(VulkanContext VkCtx, ByteBuffer data){
        int Size = data.remaining();
        stgBuffer = new VulkanBuffer(VkCtx, Size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
        long MappedMemory = stgBuffer.MapMemory(VkCtx);
        ByteBuffer buffer = MemoryUtil.memByteBuffer(MappedMemory, (int) stgBuffer.GetRequestedSize());
        buffer.put(data);
        data.flip();
        stgBuffer.UnMapMemory(VkCtx);
    }

    public void RecordTextureTransition(CommandBuffer commandBuffer){
        if(stgBuffer != null && !RecordedTransition){
            RecordedTransition = true;
            try(var MemStack = MemoryStack.stackPush()){
                VulkanUtils.ImageBarrier(MemStack, commandBuffer.GetVulkanCommandBuffer(), image.getVulkanImage(),
                        VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                        VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_2_NONE, VK_ACCESS_TRANSFER_WRITE_BIT, VK_IMAGE_ASPECT_COLOR_BIT);
                RecordCopyBuffer(MemStack, commandBuffer, stgBuffer);
                RecordGenerateMipMaps(MemStack, commandBuffer);
            }
        }
    }

    private void RecordGenerateMipMaps(MemoryStack MemStack, CommandBuffer commandBuffer){
        VkImageSubresourceRange SubResourceRange = VkImageSubresourceRange.calloc(MemStack)
                .aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
                .baseArrayLayer(0)
                .levelCount(1)
                .layerCount(1);
        VkImageMemoryBarrier2.Buffer ImageMemoryBarrierBuffer = VkImageMemoryBarrier2.calloc(1,MemStack)
                .sType$Default()
                .image(image.getVulkanImage())
                .srcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
                .dstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED)
                .subresourceRange(SubResourceRange);

        VkDependencyInfo dependencyInfo = VkDependencyInfo.calloc(MemStack)
                .sType$Default()
                .pImageMemoryBarriers(ImageMemoryBarrierBuffer);

        int MipMapWidth = Width;
        int MipMapHeight = Height;

        int MipMapLevels = image.GetMipLevels();
        for(int i = 1; i < MipMapLevels; i++){
            SubResourceRange.baseMipLevel(i - 1);
            ImageMemoryBarrierBuffer.subresourceRange(SubResourceRange)
                    .oldLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
                    .newLayout(VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL)
                    .srcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
                    .dstStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
                    .srcAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
                    .dstAccessMask(VK_ACCESS_TRANSFER_READ_BIT);
            vkCmdPipelineBarrier2(commandBuffer.GetVulkanCommandBuffer(),dependencyInfo);

            int auxi = i;
            VkOffset3D srcOffset0 = VkOffset3D.calloc(MemStack).x(0).y(0).z(0);
            VkOffset3D srcOffset1 = VkOffset3D.calloc(MemStack).x(MipMapWidth).y(MipMapHeight).z(1);
            VkOffset3D dstOffset0 = VkOffset3D.calloc(MemStack).x(0).y(0).z(0);
            VkOffset3D dstOffset1 = VkOffset3D.calloc(MemStack)
                    .x(MipMapWidth > 1 ? MipMapWidth/2 : 1).y(MipMapHeight > 1 ? MipMapHeight/2 : 1).z(1);
            VkImageBlit.Buffer BlitBuffer = VkImageBlit.calloc(1,MemStack)
                    .srcOffsets(0,srcOffset0)
                    .srcOffsets(1,srcOffset1)
                    .srcSubresource(it -> it.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
                            .mipLevel(auxi - 1)
                            .baseArrayLayer(0)
                            .layerCount(1))
                    .dstOffsets(0,dstOffset0)
                    .dstOffsets(1,dstOffset1)
                    .dstSubresource(it -> it.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
                            .mipLevel(auxi)
                            .baseArrayLayer(0)
                            .layerCount(1));
            vkCmdBlitImage(commandBuffer.GetVulkanCommandBuffer(),image.getVulkanImage(),VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                    image.getVulkanImage(),VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,BlitBuffer,VK_FILTER_LINEAR);

            ImageMemoryBarrierBuffer.oldLayout(VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL).newLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
                    .srcAccessMask(VK_ACCESS_TRANSFER_READ_BIT).dstAccessMask(VK_ACCESS_SHADER_READ_BIT);
            ImageMemoryBarrierBuffer.srcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT).dstAccessMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT);

            vkCmdPipelineBarrier2(commandBuffer.GetVulkanCommandBuffer(),dependencyInfo);
            if(MipMapWidth > 1) MipMapWidth/=2;
            if(MipMapHeight > 1) MipMapHeight/=2;
        }
        ImageMemoryBarrierBuffer.subresourceRange(it -> it.baseMipLevel(MipMapLevels - 1))
                .oldLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
                .newLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
                .srcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
                .dstAccessMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
                .srcAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
                .dstAccessMask(VK_ACCESS_SHADER_READ_BIT);

        vkCmdPipelineBarrier2(commandBuffer.GetVulkanCommandBuffer(), dependencyInfo);
    }

    private void RecordCopyBuffer(MemoryStack MemStack, CommandBuffer commandBuffer, VulkanBuffer bufferData){
        VkBufferImageCopy.Buffer Region = VkBufferImageCopy.calloc(1, MemStack)
                .bufferOffset(0)
                .bufferRowLength(0)
                .bufferImageHeight(0)
                .imageSubresource(it->
                        it.aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
                                .mipLevel(0)
                                .baseArrayLayer(0)
                                .layerCount(1)
                ).imageOffset(it->it.x(0).y(0).z(0))
                .imageExtent(it->it.width(Width).height(Height).depth(1));

        vkCmdCopyBufferToImage(commandBuffer.GetVulkanCommandBuffer(), bufferData.GetBuffer(), image.getVulkanImage(),
                VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, Region);
    }

    public void CleanUp(VulkanContext VkCtx){
        CleanUpStgBuffer(VkCtx);
        imageView.cleanup(VkCtx.GetDevice());
        image.CleanUp(VkCtx);
    }

    public void CleanUpStgBuffer(VulkanContext VkCtx){
        if(stgBuffer != null){
            stgBuffer.cleanup(VkCtx);
            stgBuffer = null;
        }
    }

    public int GetHeight(){return Height;}
    public int GetWidth(){return Width;}
    public String GetID(){return ID;}
    public ImageView GetImageView(){return imageView;}
}
