package net.vulkanmod.vulkan.shader;

import it.unimi.dsi.fastutil.objects.ObjectArrayList;
import it.unimi.dsi.fastutil.objects.ObjectListIterator;
import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.net.URL;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.LinkOption;
import java.nio.file.Path;
import java.nio.file.Paths;
import net.vulkanmod.render.shader.ShaderPackManager;
import org.lwjgl.system.MemoryStack;
import org.lwjgl.system.MemoryUtil;
import org.lwjgl.system.NativeResource;
import org.lwjgl.util.shaderc.Shaderc;
import org.lwjgl.util.shaderc.ShadercIncludeResolveI;
import org.lwjgl.util.shaderc.ShadercIncludeResult;
import org.lwjgl.util.shaderc.ShadercIncludeResultReleaseI;
import org.lwjgl.vulkan.VK12;

public class SPIRVUtils {
   private static final boolean DEBUG = true;
   private static final boolean OPTIMIZATIONS = false;
   private static long compiler;
   private static long options;
   private static final ShaderIncluder SHADER_INCLUDER = new ShaderIncluder();
   private static final ShaderReleaser SHADER_RELEASER = new ShaderReleaser();
   private static final long pUserData = 0L;
   private static ObjectArrayList<String> includePaths;

   private static void initCompiler() {
      compiler = Shaderc.shaderc_compiler_initialize();
      if (compiler == 0L) {
         throw new RuntimeException("Failed to create shader compiler");
      } else {
         options = Shaderc.shaderc_compile_options_initialize();
         if (options == 0L) {
            throw new RuntimeException("Failed to create compiler options");
         } else {
            Shaderc.shaderc_compile_options_set_generate_debug_info(options);
            Shaderc.shaderc_compile_options_set_target_env(options, 4202496, VK12.VK_API_VERSION_1_2);
            Shaderc.shaderc_compile_options_set_include_callbacks(options, SHADER_INCLUDER, SHADER_RELEASER, 0L);
            includePaths = new ObjectArrayList();
            addIncludePath("/assets/vulkanmod/shaders/include/");
         }
      }
   }

   public static void addIncludePath(String path) {
      URL url = SPIRVUtils.class.getResource(path);
      if (url != null) {
         includePaths.add(url.toExternalForm());
      }

   }

   public static SPIRV compileShader(String filename, String source, ShaderKind shaderKind) {
      if (source == null) {
         throw new NullPointerException("source for %s.%s is null".formatted(filename, shaderKind));
      } else {
         long result = Shaderc.shaderc_compile_into_spv(compiler, source, shaderKind.kind, filename, "main", options);
         if (result == 0L) {
            throw new RuntimeException("Failed to compile shader " + filename + " into SPIR-V");
         } else if (Shaderc.shaderc_result_get_compilation_status(result) != 0) {
            String errorMessage = Shaderc.shaderc_result_get_error_message(result);
            throw new RuntimeException("Failed to compile shader %s into SPIR-V:\n\t%s".formatted(filename, errorMessage));
         } else {
            return new SPIRV(result, Shaderc.shaderc_result_get_bytes(result));
         }
      }
   }

   static {
      initCompiler();
   }

   public static enum ShaderKind {
      VERTEX_SHADER(0),
      GEOMETRY_SHADER(3),
      FRAGMENT_SHADER(1),
      COMPUTE_SHADER(2);

      private final int kind;

      private ShaderKind(int kind) {
         this.kind = kind;
      }

      // $FF: synthetic method
      private static ShaderKind[] $values() {
         return new ShaderKind[]{VERTEX_SHADER, GEOMETRY_SHADER, FRAGMENT_SHADER, COMPUTE_SHADER};
      }
   }

   private static class ShaderIncluder implements ShadercIncludeResolveI {
      private static final int MAX_PATH_LENGTH = 4096;

      public long invoke(long user_data, long requested_source, int type, long requesting_source, long include_depth) {
         String requesting = MemoryUtil.memASCII(requesting_source);
         String requested = MemoryUtil.memASCII(requested_source);

         try {
            MemoryStack stack;
            label97: {
               long var26;
               label98: {
                  stack = MemoryStack.stackPush();

                  long var17;
                  try {
                     label99: {
                        InputStream shaderPackIncludeStream = ShaderPackManager.openActiveShaderResource("include/" + requested);

                        label104: {
                           try {
                              if (shaderPackIncludeStream != null) {
                                 byte[] bytes = shaderPackIncludeStream.readAllBytes();
                                 var26 = ShadercIncludeResult.malloc(stack).source_name(stack.ASCII(requested)).content(stack.bytes(bytes)).user_data(user_data).address();
                                 break label104;
                              }
                           } catch (Throwable var21) {
                              if (shaderPackIncludeStream != null) {
                                 try {
                                    shaderPackIncludeStream.close();
                                 } catch (Throwable var20) {
                                    var21.addSuppressed(var20);
                                 }
                              }

                              throw var21;
                           }

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

                           ObjectListIterator var24 = SPIRVUtils.includePaths.iterator();

                           Path path;
                           do {
                              if (!var24.hasNext()) {
                                 break label97;
                              }

                              String includePath = (String)var24.next();
                              path = Paths.get(new URI(String.format("%s%s", includePath, requested)));
                           } while(!Files.exists(path, new LinkOption[0]));

                           byte[] bytes = Files.readAllBytes(path);
                           var17 = ShadercIncludeResult.malloc(stack).source_name(stack.ASCII(requested)).content(stack.bytes(bytes)).user_data(user_data).address();
                           break label99;
                        }

                        if (shaderPackIncludeStream != null) {
                           shaderPackIncludeStream.close();
                        }
                        break label98;
                     }
                  } catch (Throwable var22) {
                     if (stack != null) {
                        try {
                           stack.close();
                        } catch (Throwable var19) {
                           var22.addSuppressed(var19);
                        }
                     }

                     throw var22;
                  }

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

                  return var17;
               }

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

               return var26;
            }

            if (stack != null) {
               stack.close();
            }
         } catch (URISyntaxException | IOException e) {
            throw new RuntimeException(e);
         }

         throw new RuntimeException(String.format("%s: Unable to find %s in include paths", requesting, requested));
      }
   }

   private static class ShaderReleaser implements ShadercIncludeResultReleaseI {
      public void invoke(long user_data, long include_result) {
      }
   }

   public static final class SPIRV implements NativeResource {
      private final long handle;
      private ByteBuffer bytecode;

      public SPIRV(long handle, ByteBuffer bytecode) {
         this.handle = handle;
         this.bytecode = bytecode;
      }

      public ByteBuffer bytecode() {
         return this.bytecode;
      }

      public void free() {
         this.bytecode = null;
      }
   }
}