package net.vulkanmod.vulkan.framebuffer;

import it.unimi.dsi.fastutil.objects.Reference2LongArrayMap;
import java.nio.LongBuffer;
import java.util.Arrays;
import net.vulkanmod.vulkan.Vulkan;
import net.vulkanmod.vulkan.memory.MemoryManager;
import net.vulkanmod.vulkan.texture.VulkanImage;
import org.apache.commons.lang3.Validate;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.vulkan.VK10;
import org.lwjgl.vulkan.VkAllocationCallbacks;
import org.lwjgl.vulkan.VkCommandBuffer;
import org.lwjgl.vulkan.VkDevice;
import org.lwjgl.vulkan.VkFramebufferCreateInfo;
import org.lwjgl.vulkan.VkRect2D;
import org.lwjgl.vulkan.VkViewport;

public class Framebuffer {
   public static final int DEFAULT_FORMAT = 37;
   protected int format;
   protected int depthFormat;
   protected int width;
   protected int height;
   protected boolean linearFiltering;
   protected boolean depthLinearFiltering;
   protected int attachmentCount;
   boolean hasColorAttachment;
   boolean hasDepthAttachment;
   private VulkanImage colorAttachment;
   protected VulkanImage depthAttachment;
   private final Reference2LongArrayMap<RenderPass> renderpassToFramebufferMap = new Reference2LongArrayMap();

   protected Framebuffer() {
   }

   public Framebuffer(Builder builder) {
      this.format = builder.format;
      this.depthFormat = builder.depthFormat;
      this.width = builder.width;
      this.height = builder.height;
      this.linearFiltering = builder.linearFiltering;
      this.depthLinearFiltering = builder.depthLinearFiltering;
      this.hasColorAttachment = builder.hasColorAttachment;
      this.hasDepthAttachment = builder.hasDepthAttachment;
      if (builder.createImages) {
         this.createImages();
      } else {
         this.colorAttachment = builder.colorAttachment;
         this.depthAttachment = builder.depthAttachment;
      }

   }

   public void createImages() {
      if (this.hasColorAttachment) {
         this.colorAttachment = VulkanImage.builder(this.width, this.height).setFormat(this.format).setUsage(20).setLinearFiltering(this.linearFiltering).setClamp(true).createVulkanImage();
      }

      if (this.hasDepthAttachment) {
         this.depthAttachment = VulkanImage.createDepthImage(this.depthFormat, this.width, this.height, 36, this.depthLinearFiltering, true);
         ++this.attachmentCount;
      }

   }

   public void resize(int newWidth, int newHeight) {
      this.width = newWidth;
      this.height = newHeight;
      this.cleanUp();
      this.createImages();
   }

   private long createFramebuffer(RenderPass renderPass) {
      MemoryStack stack = MemoryStack.stackPush();

      long var6;
      try {
         LongBuffer attachments;
         if (this.colorAttachment != null && this.depthAttachment != null) {
            attachments = stack.longs(this.colorAttachment.getImageView(), this.depthAttachment.getImageView());
         } else if (this.colorAttachment != null) {
            attachments = stack.longs(this.colorAttachment.getImageView());
         } else {
            if (this.depthAttachment == null) {
               throw new IllegalStateException();
            }

            attachments = stack.longs(this.depthAttachment.getImageView());
         }

         LongBuffer pFramebuffer = stack.mallocLong(1);
         VkFramebufferCreateInfo framebufferInfo = VkFramebufferCreateInfo.calloc(stack);
         framebufferInfo.sType$Default();
         framebufferInfo.renderPass(renderPass.getId());
         framebufferInfo.width(this.width);
         framebufferInfo.height(this.height);
         framebufferInfo.layers(1);
         framebufferInfo.pAttachments(attachments);
         if (VK10.vkCreateFramebuffer(Vulkan.getVkDevice(), framebufferInfo, (VkAllocationCallbacks)null, pFramebuffer) != 0) {
            throw new RuntimeException("Failed to create framebuffer");
         }

         var6 = pFramebuffer.get(0);
      } catch (Throwable var9) {
         if (stack != null) {
            try {
               stack.close();
            } catch (Throwable var8) {
               var9.addSuppressed(var8);
            }
         }

         throw var9;
      }

      if (stack != null) {
         stack.close();
      }

      return var6;
   }

   public void beginRenderPass(VkCommandBuffer commandBuffer, RenderPass renderPass, MemoryStack stack) {
      long framebufferId = this.getFramebufferId(renderPass);
      renderPass.beginRenderPass(commandBuffer, framebufferId, stack);
   }

   protected long getFramebufferId(RenderPass renderPass) {
      return this.renderpassToFramebufferMap.computeIfAbsent(renderPass, (renderPass1) -> this.createFramebuffer(renderPass));
   }

   public VkViewport.Buffer viewport(MemoryStack stack) {
      VkViewport.Buffer viewport = VkViewport.malloc(1, stack);
      viewport.x(0.0F);
      viewport.y((float)this.height);
      viewport.width((float)this.width);
      viewport.height((float)(-this.height));
      viewport.minDepth(0.0F);
      viewport.maxDepth(1.0F);
      return viewport;
   }

   public VkRect2D.Buffer scissor(MemoryStack stack) {
      VkRect2D.Buffer scissor = VkRect2D.malloc(1, stack);
      scissor.offset().set(0, 0);
      scissor.extent().set(this.width, this.height);
      return scissor;
   }

   public void cleanUp() {
      this.cleanUp(true);
   }

   public void cleanUp(boolean cleanImages) {
      if (cleanImages) {
         if (this.colorAttachment != null) {
            this.colorAttachment.free();
         }

         if (this.depthAttachment != null) {
            this.depthAttachment.free();
         }
      }

      VkDevice device = Vulkan.getVkDevice();
      long[] ids = this.renderpassToFramebufferMap.values().toLongArray();
      MemoryManager.getInstance().addFrameOp(() -> Arrays.stream(ids).forEach((id) -> VK10.vkDestroyFramebuffer(device, id, (VkAllocationCallbacks)null)));
      this.renderpassToFramebufferMap.clear();
   }

   public long getDepthImageView() {
      return this.depthAttachment.getImageView();
   }

   public VulkanImage getDepthAttachment() {
      return this.depthAttachment;
   }

   public VulkanImage getColorAttachment() {
      return this.colorAttachment;
   }

   public int getWidth() {
      return this.width;
   }

   public int getHeight() {
      return this.height;
   }

   public int getFormat() {
      return this.format;
   }

   public int getDepthFormat() {
      return this.depthFormat;
   }

   public static Builder builder(int width, int height, int colorAttachments, boolean hasDepthAttachment) {
      return new Builder(width, height, colorAttachments, hasDepthAttachment);
   }

   public static Builder builder(VulkanImage colorAttachment, VulkanImage depthAttachment) {
      return new Builder(colorAttachment, depthAttachment);
   }

   public static class Builder {
      final boolean createImages;
      final int width;
      final int height;
      int format;
      int depthFormat;
      VulkanImage colorAttachment;
      VulkanImage depthAttachment;
      boolean hasColorAttachment;
      boolean hasDepthAttachment;
      boolean linearFiltering;
      boolean depthLinearFiltering;

      public Builder(int width, int height, int colorAttachments, boolean hasDepthAttachment) {
         Validate.isTrue(colorAttachments > 0 || hasDepthAttachment, "At least 1 attachment needed", new Object[0]);
         Validate.isTrue(colorAttachments <= 1, "Not supported", new Object[0]);
         this.createImages = true;
         this.format = 37;
         this.depthFormat = Vulkan.getDefaultDepthFormat();
         this.linearFiltering = true;
         this.depthLinearFiltering = false;
         this.width = width;
         this.height = height;
         this.hasColorAttachment = colorAttachments == 1;
         this.hasDepthAttachment = hasDepthAttachment;
      }

      public Builder(VulkanImage colorAttachment, VulkanImage depthAttachment) {
         this.createImages = false;
         this.colorAttachment = colorAttachment;
         this.depthAttachment = depthAttachment;
         this.format = colorAttachment.format;
         this.width = colorAttachment.width;
         this.height = colorAttachment.height;
         this.hasColorAttachment = true;
         this.hasDepthAttachment = depthAttachment != null;
         this.depthFormat = this.hasDepthAttachment ? depthAttachment.format : 0;
         this.linearFiltering = true;
         this.depthLinearFiltering = false;
      }

      public Framebuffer build() {
         return new Framebuffer(this);
      }

      public Builder setFormat(int format) {
         this.format = format;
         return this;
      }

      public Builder setLinearFiltering(boolean b) {
         this.linearFiltering = b;
         return this;
      }

      public Builder setDepthLinearFiltering(boolean b) {
         this.depthLinearFiltering = b;
         return this;
      }

      public Builder setDepthFormat(int format) {
         this.depthFormat = format;
         return this;
      }
   }
}