package net.vulkanmod.render;

import com.google.gson.JsonObject;
import com.mojang.blaze3d.vertex.VertexFormat;
import java.util.function.Function;
import net.minecraft.class_1921;
import net.minecraft.class_290;
import net.vulkanmod.render.chunk.build.thread.ThreadBuilderPack;
import net.vulkanmod.render.shader.ShaderLoadUtil;
import net.vulkanmod.render.vertex.CustomVertexFormat;
import net.vulkanmod.render.vertex.TerrainRenderType;
import net.vulkanmod.vulkan.shader.GraphicsPipeline;
import net.vulkanmod.vulkan.shader.Pipeline;
import net.vulkanmod.vulkan.shader.descriptor.UBO;

public abstract class PipelineManager {
   public static VertexFormat terrainVertexFormat;
   static GraphicsPipeline terrainShader;
   static GraphicsPipeline terrainShaderEarlyZ;
   static GraphicsPipeline waterShader;
   static GraphicsPipeline fastBlitPipeline;
   static GraphicsPipeline cloudsPipeline;
   static GraphicsPipeline shadowTerrainShader;
   static GraphicsPipeline shadowEntityShader;
   private static Function<TerrainRenderType, GraphicsPipeline> shaderGetter;

   public static void setTerrainVertexFormat(VertexFormat format) {
      terrainVertexFormat = format;
   }

   public static void init() {
      setTerrainVertexFormat(CustomVertexFormat.COMPRESSED_TERRAIN);
      createBasicPipelines();
      setDefaultShader();
      ThreadBuilderPack.defaultTerrainBuilderConstructor();
   }

   public static void setDefaultShader() {
      setShaderGetter((renderType) -> renderType == TerrainRenderType.TRANSLUCENT ? waterShader : terrainShader);
   }

   private static void createBasicPipelines() {
      terrainShaderEarlyZ = createPipeline("terrain_earlyZ", terrainVertexFormat);
      terrainShader = createPipeline("terrain", terrainVertexFormat);
      waterShader = createPipeline("terrain_water", terrainVertexFormat);
      fastBlitPipeline = createPipeline("blit", CustomVertexFormat.NONE);
      cloudsPipeline = createPipeline("clouds", class_290.field_1576);
      shadowTerrainShader = createPipeline("shadow", terrainVertexFormat);
      shadowEntityShader = createPipeline("entity_shadow", class_290.field_1592);
   }

   private static GraphicsPipeline createPipeline(String configName, VertexFormat vertexFormat) {
      Pipeline.Builder pipelineBuilder = new Pipeline.Builder(vertexFormat, configName);
      String path = ShaderLoadUtil.resolveShaderPath("basic");
      JsonObject config = ShaderLoadUtil.getJsonConfig(path, configName);
      if (config == null) {
         path = ShaderLoadUtil.resolveBuiltinShaderPath("basic");
         config = ShaderLoadUtil.getJsonConfig(path, configName);
      }

      if (config == null) {
         throw new IllegalStateException("Unable to find shader config for pipeline " + configName);
      } else {
         pipelineBuilder.parseBindings(config);
         ShaderLoadUtil.loadShaders(pipelineBuilder, config, configName, path);
         GraphicsPipeline pipeline = pipelineBuilder.createGraphicsPipeline();

         for(UBO buffer : pipeline.getBuffers()) {
            buffer.setUseGlobalBuffer(true);
         }

         return pipeline;
      }
   }

   public static void reloadPipelines() {
      destroyPipelines();
      createBasicPipelines();
      setDefaultShader();
      ThreadBuilderPack.defaultTerrainBuilderConstructor();
   }

   public static GraphicsPipeline getTerrainShader(TerrainRenderType renderType) {
      return (GraphicsPipeline)shaderGetter.apply(renderType);
   }

   public static void setShaderGetter(Function<TerrainRenderType, GraphicsPipeline> consumer) {
      shaderGetter = consumer;
   }

   public static GraphicsPipeline getTerrainDirectShader(class_1921 renderType) {
      return terrainShader;
   }

   public static GraphicsPipeline getTerrainIndirectShader(class_1921 renderType) {
      return terrainShaderEarlyZ;
   }

   public static GraphicsPipeline getFastBlitPipeline() {
      return fastBlitPipeline;
   }

   public static GraphicsPipeline getCloudsPipeline() {
      return cloudsPipeline;
   }

   public static GraphicsPipeline getShadowTerrainShader() {
      return shadowTerrainShader;
   }

   public static GraphicsPipeline getShadowEntityShader() {
      return shadowEntityShader;
   }

   public static void destroyPipelines() {
      if (terrainShaderEarlyZ != null) {
         terrainShaderEarlyZ.cleanUp();
      }

      if (terrainShader != null) {
         terrainShader.cleanUp();
      }

      if (waterShader != null) {
         waterShader.cleanUp();
      }

      if (fastBlitPipeline != null) {
         fastBlitPipeline.cleanUp();
      }

      if (cloudsPipeline != null) {
         cloudsPipeline.cleanUp();
      }

      if (shadowTerrainShader != null) {
         shadowTerrainShader.cleanUp();
      }

      if (shadowEntityShader != null) {
         shadowEntityShader.cleanUp();
      }

   }
}