package net.vulkanmod.render.chunk.buffer;

import java.nio.ByteBuffer;
import java.util.EnumMap;
import java.util.Iterator;
import net.minecraft.class_243;
import net.vulkanmod.Initializer;
import net.vulkanmod.render.PipelineManager;
import net.vulkanmod.render.chunk.RenderSection;
import net.vulkanmod.render.chunk.build.UploadBuffer;
import net.vulkanmod.render.chunk.cull.QuadFacing;
import net.vulkanmod.render.chunk.util.StaticQueue;
import net.vulkanmod.render.vertex.CustomVertexFormat;
import net.vulkanmod.render.vertex.TerrainRenderType;
import net.vulkanmod.vulkan.Renderer;
import net.vulkanmod.vulkan.memory.buffer.IndirectBuffer;
import net.vulkanmod.vulkan.shader.Pipeline;
import org.joml.Vector3i;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.vulkan.VK10;
import org.lwjgl.vulkan.VkCommandBuffer;

public class DrawBuffers {
   public static final int VERTEX_SIZE;
   public static final int INDEX_SIZE = 2;
   public static final int UNDEFINED_FACING_IDX;
   public static final float POS_OFFSET;
   private static final int CMD_STRIDE = 32;
   private static final long cmdBufferPtr;
   private final int index;
   private final Vector3i origin;
   private final int minHeight;
   private boolean allocated = false;
   AreaBuffer indexBuffer;
   private final EnumMap<TerrainRenderType, AreaBuffer> vertexBuffers = new EnumMap(TerrainRenderType.class);
   long drawParamsPtr;
   final int[] sectionIndices = new int[512];
   final int[] masks = new int[512];

   public DrawBuffers(int index, Vector3i origin, int minHeight) {
      this.index = index;
      this.origin = origin;
      this.minHeight = minHeight;
      this.drawParamsPtr = DrawParametersBuffer.allocateBuffer();
   }

   public void upload(RenderSection section, UploadBuffer buffer, TerrainRenderType renderType) {
      ByteBuffer[] vertexBuffers = buffer.getVertexBuffers();
      if (buffer.indexOnly) {
         long paramsPtr = DrawParametersBuffer.getParamsPtr(this.drawParamsPtr, section.inAreaIndex, renderType.ordinal(), QuadFacing.UNDEFINED.ordinal());
         int firstIndex = DrawParametersBuffer.getFirstIndex(paramsPtr);
         int indexCount = DrawParametersBuffer.getIndexCount(paramsPtr);
         int oldOffset = indexCount > 0 ? firstIndex : -1;
         AreaBuffer.Segment segment = this.indexBuffer.upload(buffer.getIndexBuffer(), oldOffset, paramsPtr);
         firstIndex = segment.offset / 2;
         DrawParametersBuffer.setFirstIndex(paramsPtr, firstIndex);
         buffer.release();
      } else {
         int oldOffset = -1;
         int size = 0;

         for(int i = 0; i < QuadFacing.COUNT; ++i) {
            long paramPtr = DrawParametersBuffer.getParamsPtr(this.drawParamsPtr, section.inAreaIndex, renderType.ordinal(), i);
            int vertexOffset = DrawParametersBuffer.getVertexOffset(paramPtr);
            if (oldOffset == -1) {
               oldOffset = vertexOffset;
            }

            ByteBuffer vertexBuffer = vertexBuffers[i];
            if (vertexBuffer != null) {
               size += vertexBuffer.remaining();
            }
         }

         AreaBuffer areaBuffer = null;
         AreaBuffer.Segment segment = null;
         boolean doUpload = false;
         if (size > 0) {
            areaBuffer = this.getAreaBufferOrAlloc(renderType);
            areaBuffer.freeSegment(oldOffset);
            segment = areaBuffer.allocateSegment(size);
            doUpload = true;
         }

         int baseInstance = this.encodeSectionOffset(section.xOffset(), section.yOffset(), section.zOffset());
         int offset = 0;

         for(int i = 0; i < QuadFacing.COUNT; ++i) {
            long paramPtr = DrawParametersBuffer.getParamsPtr(this.drawParamsPtr, section.inAreaIndex, renderType.ordinal(), i);
            int vertexOffset = -1;
            int firstIndex = 0;
            int indexCount = 0;
            ByteBuffer vertexBuffer = vertexBuffers[i];
            int vertexCount = 0;
            if (vertexBuffer != null && doUpload) {
               areaBuffer.upload(segment, vertexBuffer, offset);
               vertexOffset = (segment.offset + offset) / VERTEX_SIZE;
               offset += vertexBuffer.remaining();
               vertexCount = vertexBuffer.limit() / VERTEX_SIZE;
               indexCount = vertexCount * 6 / 4;
            }

            if (i == QuadFacing.UNDEFINED.ordinal() && !buffer.autoIndices) {
               if (this.indexBuffer == null) {
                  this.indexBuffer = new AreaBuffer(AreaBuffer.Usage.INDEX, 60000, 2);
               }

               oldOffset = DrawParametersBuffer.getIndexCount(paramPtr) > 0 ? DrawParametersBuffer.getFirstIndex(paramPtr) : -1;
               AreaBuffer.Segment ibSegment = this.indexBuffer.upload(buffer.getIndexBuffer(), oldOffset, paramPtr);
               firstIndex = ibSegment.offset / 2;
            } else {
               Renderer.getDrawer().getQuadsIndexBuffer().checkCapacity(vertexCount);
            }

            DrawParametersBuffer.setIndexCount(paramPtr, indexCount);
            DrawParametersBuffer.setFirstIndex(paramPtr, firstIndex);
            DrawParametersBuffer.setVertexOffset(paramPtr, vertexOffset);
            DrawParametersBuffer.setBaseInstance(paramPtr, baseInstance);
         }

         buffer.release();
      }
   }

   private AreaBuffer getAreaBufferOrAlloc(TerrainRenderType renderType) {
      this.allocated = true;
      int var10000;
      switch (renderType) {
         case SOLID:
         case CUTOUT:
            var10000 = 100000;
            break;
         case CUTOUT_MIPPED:
            var10000 = 250000;
            break;
         case TRANSLUCENT:
         case TRIPWIRE:
            var10000 = 60000;
            break;
         default:
            throw new MatchException((String)null, (Throwable)null);
      }

      int initialSize = var10000;
      return (AreaBuffer)this.vertexBuffers.computeIfAbsent(renderType, (renderType1) -> new AreaBuffer(AreaBuffer.Usage.VERTEX, initialSize, VERTEX_SIZE));
   }

   public AreaBuffer getAreaBuffer(TerrainRenderType r) {
      return (AreaBuffer)this.vertexBuffers.get(r);
   }

   private boolean hasRenderType(TerrainRenderType r) {
      return this.vertexBuffers.containsKey(r);
   }

   private int encodeSectionOffset(int xOffset, int yOffset, int zOffset) {
      int xOffset1 = xOffset & 127;
      int zOffset1 = zOffset & 127;
      int yOffset1 = yOffset - this.minHeight & 127;
      return yOffset1 << 16 | zOffset1 << 8 | xOffset1;
   }

   private void updateChunkAreaOrigin(VkCommandBuffer commandBuffer, Pipeline pipeline, double camX, double camY, double camZ, MemoryStack stack) {
      float xOffset = (float)((double)((float)this.origin.x + POS_OFFSET) - camX);
      float yOffset = (float)((double)((float)this.origin.y + POS_OFFSET) - camY);
      float zOffset = (float)((double)((float)this.origin.z + POS_OFFSET) - camZ);
      ByteBuffer byteBuffer = stack.malloc(12);
      byteBuffer.putFloat(0, xOffset);
      byteBuffer.putFloat(4, yOffset);
      byteBuffer.putFloat(8, zOffset);
      VK10.vkCmdPushConstants(commandBuffer, pipeline.getLayout(), 1, 0, byteBuffer);
   }

   public void buildDrawBatchesIndirect(class_243 cameraPos, IndirectBuffer indirectBuffer, StaticQueue<RenderSection> queue, TerrainRenderType terrainRenderType) {
      long bufferPtr = cmdBufferPtr;
      boolean isTranslucent = terrainRenderType == TerrainRenderType.TRANSLUCENT;
      boolean backFaceCulling = Initializer.CONFIG.backFaceCulling && !isTranslucent;
      int drawCount = 0;
      long drawParamsBasePtr = this.drawParamsPtr + (long)(terrainRenderType.ordinal() * 512 * 7) * 16L;
      long facingsStride = 112L;
      int count = 0;
      if (backFaceCulling) {
         for(Iterator<RenderSection> iterator = queue.iterator(isTranslucent); iterator.hasNext(); ++count) {
            RenderSection section = (RenderSection)iterator.next();
            this.sectionIndices[count] = section.inAreaIndex;
            this.masks[count] = this.getMask(cameraPos, section);
         }

         long ptr = bufferPtr;

         for(int j = 0; j < count; ++j) {
            int sectionIdx = this.sectionIndices[j];
            int mask = this.masks[j];
            long drawParamsBasePtr2 = drawParamsBasePtr + (long)sectionIdx * 112L;
            int indexCount = 0;
            int firstIndex = 0;
            int vertexOffset = 0;
            int baseInstance = 0;

            for(int i = 0; i < QuadFacing.COUNT; ++i) {
               if ((mask & 1 << i) == 0) {
                  drawParamsBasePtr2 += 16L;
                  if (indexCount > 0) {
                     MemoryUtil.memPutInt(ptr, indexCount);
                     MemoryUtil.memPutInt(ptr + 4L, 1);
                     MemoryUtil.memPutInt(ptr + 8L, firstIndex);
                     MemoryUtil.memPutInt(ptr + 12L, vertexOffset);
                     MemoryUtil.memPutInt(ptr + 16L, baseInstance);
                     ptr += 32L;
                     ++drawCount;
                  }

                  indexCount = 0;
                  firstIndex = 0;
                  vertexOffset = 0;
                  baseInstance = 0;
               } else {
                  int indexCount_i = DrawParametersBuffer.getIndexCount(drawParamsBasePtr2);
                  int firstIndex_i = DrawParametersBuffer.getFirstIndex(drawParamsBasePtr2);
                  int vertexOffset_i = DrawParametersBuffer.getVertexOffset(drawParamsBasePtr2);
                  int baseInstance_i = DrawParametersBuffer.getBaseInstance(drawParamsBasePtr2);
                  if (indexCount == 0) {
                     indexCount = indexCount_i;
                     firstIndex = firstIndex_i;
                     vertexOffset = vertexOffset_i;
                     baseInstance = baseInstance_i;
                  } else {
                     indexCount += indexCount_i;
                  }

                  drawParamsBasePtr2 += 16L;
               }
            }

            if (indexCount > 0) {
               MemoryUtil.memPutInt(ptr, indexCount);
               MemoryUtil.memPutInt(ptr + 4L, 1);
               MemoryUtil.memPutInt(ptr + 8L, firstIndex);
               MemoryUtil.memPutInt(ptr + 12L, vertexOffset);
               MemoryUtil.memPutInt(ptr + 16L, baseInstance);
               ptr += 32L;
               ++drawCount;
            }
         }
      } else {
         for(Iterator<RenderSection> iterator = queue.iterator(isTranslucent); iterator.hasNext(); ++count) {
            RenderSection section = (RenderSection)iterator.next();
            this.sectionIndices[count] = section.inAreaIndex;
         }

         long facingOffset = (long)UNDEFINED_FACING_IDX * 16L;
         drawParamsBasePtr += facingOffset;
         long ptr = bufferPtr;

         for(int i = 0; i < count; ++i) {
            int sectionIdx = this.sectionIndices[i];
            long drawParamsPtr = drawParamsBasePtr + (long)sectionIdx * 112L;
            int indexCount = DrawParametersBuffer.getIndexCount(drawParamsPtr);
            int firstIndex = DrawParametersBuffer.getFirstIndex(drawParamsPtr);
            int vertexOffset = DrawParametersBuffer.getVertexOffset(drawParamsPtr);
            int baseInstance = DrawParametersBuffer.getBaseInstance(drawParamsPtr);
            if (indexCount > 0) {
               MemoryUtil.memPutInt(ptr, indexCount);
               MemoryUtil.memPutInt(ptr + 4L, 1);
               MemoryUtil.memPutInt(ptr + 8L, firstIndex);
               MemoryUtil.memPutInt(ptr + 12L, vertexOffset);
               MemoryUtil.memPutInt(ptr + 16L, baseInstance);
               ptr += 32L;
               ++drawCount;
            }
         }
      }

      if (drawCount != 0) {
         ByteBuffer byteBuffer = MemoryUtil.memByteBuffer(cmdBufferPtr, queue.size() * QuadFacing.COUNT * 32);
         indirectBuffer.recordCopyCmd(byteBuffer.position(0));
         VK10.vkCmdDrawIndexedIndirect(Renderer.getCommandBuffer(), indirectBuffer.getId(), indirectBuffer.getOffset(), drawCount, 32);
      }
   }

   public void buildDrawBatchesDirect(class_243 cameraPos, StaticQueue<RenderSection> queue, TerrainRenderType terrainRenderType) {
      boolean isTranslucent = terrainRenderType == TerrainRenderType.TRANSLUCENT;
      boolean backFaceCulling = Initializer.CONFIG.backFaceCulling && !isTranslucent;
      VkCommandBuffer commandBuffer = Renderer.getCommandBuffer();
      long drawParamsBasePtr = this.drawParamsPtr + (long)(terrainRenderType.ordinal() * 512 * 7) * 16L;
      long facingsStride = 112L;
      int count = 0;
      if (backFaceCulling) {
         for(Iterator<RenderSection> iterator = queue.iterator(isTranslucent); iterator.hasNext(); ++count) {
            RenderSection section = (RenderSection)iterator.next();
            this.sectionIndices[count] = section.inAreaIndex;
            this.masks[count] = this.getMask(cameraPos, section);
         }

         for(int j = 0; j < count; ++j) {
            int sectionIdx = this.sectionIndices[j];
            int mask = this.masks[j];
            long drawParamsBasePtr2 = drawParamsBasePtr + (long)sectionIdx * 112L;
            int indexCount = 0;
            int firstIndex = 0;
            int vertexOffset = 0;
            int baseInstance = 0;

            for(int i = 0; i < QuadFacing.COUNT; ++i) {
               if ((mask & 1 << i) == 0) {
                  drawParamsBasePtr2 += 16L;
                  if (indexCount > 0) {
                     VK10.vkCmdDrawIndexed(commandBuffer, indexCount, 1, firstIndex, vertexOffset, baseInstance);
                  }

                  indexCount = 0;
                  firstIndex = 0;
                  vertexOffset = 0;
                  baseInstance = 0;
               } else {
                  int indexCount_i = DrawParametersBuffer.getIndexCount(drawParamsBasePtr2);
                  int firstIndex_i = DrawParametersBuffer.getFirstIndex(drawParamsBasePtr2);
                  int vertexOffset_i = DrawParametersBuffer.getVertexOffset(drawParamsBasePtr2);
                  int baseInstance_i = DrawParametersBuffer.getBaseInstance(drawParamsBasePtr2);
                  if (indexCount == 0) {
                     indexCount = indexCount_i;
                     firstIndex = firstIndex_i;
                     vertexOffset = vertexOffset_i;
                     baseInstance = baseInstance_i;
                  } else {
                     indexCount += indexCount_i;
                  }

                  drawParamsBasePtr2 += 16L;
               }
            }

            if (indexCount > 0) {
               VK10.vkCmdDrawIndexed(commandBuffer, indexCount, 1, firstIndex, vertexOffset, baseInstance);
            }
         }
      } else {
         long facingOffset = (long)UNDEFINED_FACING_IDX * 16L;
         drawParamsBasePtr += facingOffset;

         for(Iterator<RenderSection> iterator = queue.iterator(isTranslucent); iterator.hasNext(); ++count) {
            RenderSection section = (RenderSection)iterator.next();
            this.sectionIndices[count] = section.inAreaIndex;
         }

         for(int i = 0; i < count; ++i) {
            int sectionIdx = this.sectionIndices[i];
            long drawParamsPtr = drawParamsBasePtr + (long)sectionIdx * 112L;
            int indexCount = DrawParametersBuffer.getIndexCount(drawParamsPtr);
            int firstIndex = DrawParametersBuffer.getFirstIndex(drawParamsPtr);
            int vertexOffset = DrawParametersBuffer.getVertexOffset(drawParamsPtr);
            int baseInstance = DrawParametersBuffer.getBaseInstance(drawParamsPtr);
            if (indexCount > 0) {
               VK10.vkCmdDrawIndexed(commandBuffer, indexCount, 1, firstIndex, vertexOffset, baseInstance);
            }
         }
      }

   }

   private int getMask(class_243 camera, RenderSection section) {
      int secX = section.xOffset;
      int secY = section.yOffset;
      int secZ = section.zOffset;
      int mask = 1 << UNDEFINED_FACING_IDX;
      mask |= camera.field_1352 - (double)secX >= (double)0.0F ? 1 << QuadFacing.X_POS.ordinal() : 0;
      mask |= camera.field_1351 - (double)secY >= (double)0.0F ? 1 << QuadFacing.Y_POS.ordinal() : 0;
      mask |= camera.field_1350 - (double)secZ >= (double)0.0F ? 1 << QuadFacing.Z_POS.ordinal() : 0;
      mask |= camera.field_1352 - (double)(secX + 16) < (double)0.0F ? 1 << QuadFacing.X_NEG.ordinal() : 0;
      mask |= camera.field_1351 - (double)(secY + 16) < (double)0.0F ? 1 << QuadFacing.Y_NEG.ordinal() : 0;
      mask |= camera.field_1350 - (double)(secZ + 16) < (double)0.0F ? 1 << QuadFacing.Z_NEG.ordinal() : 0;
      return mask;
   }

   public void bindBuffers(VkCommandBuffer commandBuffer, Pipeline pipeline, TerrainRenderType terrainRenderType, double camX, double camY, double camZ) {
      MemoryStack stack = MemoryStack.stackPush();

      try {
         AreaBuffer vertexBuffer = this.getAreaBuffer(terrainRenderType);
         VK10.nvkCmdBindVertexBuffers(commandBuffer, 0, 1, stack.npointer(vertexBuffer.getId()), stack.npointer(0L));
         this.updateChunkAreaOrigin(commandBuffer, pipeline, camX, camY, camZ, stack);
      } catch (Throwable var14) {
         if (stack != null) {
            try {
               stack.close();
            } catch (Throwable var13) {
               var14.addSuppressed(var13);
            }
         }

         throw var14;
      }

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

      if (terrainRenderType == TerrainRenderType.TRANSLUCENT && this.indexBuffer != null) {
         VK10.vkCmdBindIndexBuffer(commandBuffer, this.indexBuffer.getId(), 0L, 0);
      }

   }

   public void releaseBuffers() {
      if (this.allocated) {
         this.vertexBuffers.values().forEach(AreaBuffer::freeBuffer);
         this.vertexBuffers.clear();
         if (this.indexBuffer != null) {
            this.indexBuffer.freeBuffer();
         }

         this.indexBuffer = null;
         this.allocated = false;
      }
   }

   public void free() {
      this.releaseBuffers();
      DrawParametersBuffer.freeBuffer(this.drawParamsPtr);
   }

   public boolean isAllocated() {
      return !this.vertexBuffers.isEmpty();
   }

   public EnumMap<TerrainRenderType, AreaBuffer> getVertexBuffers() {
      return this.vertexBuffers;
   }

   public AreaBuffer getIndexBuffer() {
      return this.indexBuffer;
   }

   public long getDrawParamsPtr() {
      return this.drawParamsPtr;
   }

   static {
      VERTEX_SIZE = PipelineManager.terrainVertexFormat.getVertexSize();
      UNDEFINED_FACING_IDX = QuadFacing.UNDEFINED.ordinal();
      POS_OFFSET = CustomVertexFormat.getPositionOffset();
      cmdBufferPtr = MemoryUtil.nmemAlignedAlloc(32L, 512L * (long)QuadFacing.COUNT * 32L);
   }
}