#include <glslang/Public/ShaderLang.h>
#include <glslang/SPIRV/GlslangToSpv.h>
#include <iostream>

#include "types.h"

const TBuiltInResource DefaultTBuiltInResource = {
  .maxLights                                 = 32,
  .maxClipPlanes                             = 6,
  .maxTextureUnits                           = 32,
  .maxTextureCoords                          = 32,
  .maxVertexAttribs                          = 64,
  .maxVertexUniformComponents                = 4096,
  .maxVaryingFloats                          = 64,
  .maxVertexTextureImageUnits                = 32,
  .maxCombinedTextureImageUnits              = 80,
  .maxTextureImageUnits                      = 32,
  .maxFragmentUniformComponents              = 4096,
  .maxDrawBuffers                            = 32,
  .maxVertexUniformVectors                   = 128,
  .maxVaryingVectors                         = 8,
  .maxFragmentUniformVectors                 = 16,
  .maxVertexOutputVectors                    = 16,
  .maxFragmentInputVectors                   = 15,
  .minProgramTexelOffset                     = -8,
  .maxProgramTexelOffset                     = 7,
  .maxClipDistances                          = 8,
  .maxComputeWorkGroupCountX                 = 65535,
  .maxComputeWorkGroupCountY                 = 65535,
  .maxComputeWorkGroupCountZ                 = 65535,
  .maxComputeWorkGroupSizeX                  = 1024,
  .maxComputeWorkGroupSizeY                  = 1024,
  .maxComputeWorkGroupSizeZ                  = 64,
  .maxComputeUniformComponents               = 1024,
  .maxComputeTextureImageUnits               = 16,
  .maxComputeImageUniforms                   = 8,
  .maxComputeAtomicCounters                  = 8,
  .maxComputeAtomicCounterBuffers            = 1,
  .maxVaryingComponents                      = 60,
  .maxVertexOutputComponents                 = 64,
  .maxGeometryInputComponents                = 64,
  .maxGeometryOutputComponents               = 128,
  .maxFragmentInputComponents                = 128,
  .maxImageUnits                             = 8,
  .maxCombinedImageUnitsAndFragmentOutputs   = 8,
  .maxCombinedShaderOutputResources          = 8,
  .maxImageSamples                           = 0,
  .maxVertexImageUniforms                    = 0,
  .maxTessControlImageUniforms               = 0,
  .maxTessEvaluationImageUniforms            = 0,
  .maxGeometryImageUniforms                  = 0,
  .maxFragmentImageUniforms                  = 8,
  .maxCombinedImageUniforms                  = 8,
  .maxGeometryTextureImageUnits              = 16,
  .maxGeometryOutputVertices                 = 256,
  .maxGeometryTotalOutputComponents          = 1024,
  .maxGeometryUniformComponents              = 1024,
  .maxGeometryVaryingComponents              = 64,
  .maxTessControlInputComponents             = 128,
  .maxTessControlOutputComponents            = 128,
  .maxTessControlTextureImageUnits           = 16,
  .maxTessControlUniformComponents           = 1024,
  .maxTessControlTotalOutputComponents       = 4096,
  .maxTessEvaluationInputComponents          = 128,
  .maxTessEvaluationOutputComponents         = 128,
  .maxTessEvaluationTextureImageUnits        = 16,
  .maxTessEvaluationUniformComponents        = 1024,
  .maxTessPatchComponents                    = 120,
  .maxPatchVertices                          = 32,
  .maxTessGenLevel                           = 64,
  .maxViewports                              = 16,
  .maxVertexAtomicCounters                   = 0,
  .maxTessControlAtomicCounters              = 0,
  .maxTessEvaluationAtomicCounters           = 0,
  .maxGeometryAtomicCounters                 = 0,
  .maxFragmentAtomicCounters                 = 8,
  .maxCombinedAtomicCounters                 = 8,
  .maxAtomicCounterBindings                  = 1,
  .maxVertexAtomicCounterBuffers             = 0,
  .maxTessControlAtomicCounterBuffers        = 0,
  .maxTessEvaluationAtomicCounterBuffers     = 0,
  .maxGeometryAtomicCounterBuffers           = 0,
  .maxFragmentAtomicCounterBuffers           = 1,
  .maxCombinedAtomicCounterBuffers           = 1,
  .maxAtomicCounterBufferSize                = 16384,
  .maxTransformFeedbackBuffers               = 4,
  .maxTransformFeedbackInterleavedComponents = 64,
  .maxCullDistances                          = 8,
  .maxCombinedClipAndCullDistances           = 8,
  .maxSamples                                = 4,
  .maxMeshOutputVerticesNV                   = 256,
  .maxMeshOutputPrimitivesNV                 = 512,
  .maxMeshWorkGroupSizeX_NV                  = 32,
  .maxMeshWorkGroupSizeY_NV                  = 1,
  .maxMeshWorkGroupSizeZ_NV                  = 1,
  .maxTaskWorkGroupSizeX_NV                  = 32,
  .maxTaskWorkGroupSizeY_NV                  = 1,
  .maxTaskWorkGroupSizeZ_NV                  = 1,
  .maxMeshViewCountNV                        = 4,
  .maxMeshOutputVerticesEXT                  = 256,
  .maxMeshOutputPrimitivesEXT                = 512,
  .maxMeshWorkGroupSizeX_EXT                 = 32,
  .maxMeshWorkGroupSizeY_EXT                 = 1,
  .maxMeshWorkGroupSizeZ_EXT                 = 1,
  .maxTaskWorkGroupSizeX_EXT                 = 32,
  .maxTaskWorkGroupSizeY_EXT                 = 1,
  .maxTaskWorkGroupSizeZ_EXT                 = 1,
  .maxMeshViewCountEXT                       = 4,
  .maxDualSourceDrawBuffersEXT               = 1,
  .limits                                    = {
    .nonInductiveForLoops                 = true,
    .whileLoops                           = true,
    .doWhileLoops                         = true,
    .generalUniformIndexing               = true,
    .generalAttributeMatrixVectorIndexing = true,
    .generalVaryingIndexing               = true,
    .generalSamplerIndexing               = true,
    .generalVariableIndexing              = true,
    .generalConstantMatrixVectorIndexing  = true,
  },
};

constexpr const char* vertShaderSrc = R"glsl(
  #version 450

  layout(binding = 0) uniform UniformBufferObject {
    mat4 model;
    mat4 view;
    mat4 proj;
  } ubo;

  layout(location = 0) in vec3 inPosition;
  layout(location = 1) in vec3 inColor;
  layout(location = 2) in vec2 inTexCoord;

  layout(location = 0) out vec3 fragColor;
  layout(location = 1) out vec2 fragTexCoord;

  void main() {
    gl_Position = ubo.proj * ubo.view * ubo.model * vec4(inPosition, 1.0);
    fragColor = inColor;
    fragTexCoord = inTexCoord;
  }
)glsl";

constexpr const char* fragShaderSrc = R"glsl(
  #version 450

  layout(binding = 1) uniform sampler2D texSampler;

  layout(location = 0) in vec3 fragColor;
  layout(location = 1) in vec2 fragTexCoord;

  layout(location = 0) out vec4 outColor;

  void main() {
    outColor = texture(texSampler, fragTexCoord);
  }
)glsl";

inline void InitGlslang() { glslang::InitializeProcess(); }

inline void CleanupGlslang() { glslang::FinalizeProcess(); }

inline fn CompileShader(const char* source, EShLanguage shaderType) -> std::vector<uint32_t> {
  glslang::TShader shader(shaderType);
  shader.setStrings(&source, 1);

  // Set shader environment details (language version, profile)
  shader.setEnvInput(glslang::EShSourceGlsl, shaderType, glslang::EShClientVulkan, 450);
  shader.setEnvClient(glslang::EShClientVulkan, glslang::EShTargetVulkan_1_2);
  shader.setEnvTarget(glslang::EShTargetSpv, glslang::EShTargetSpv_1_3);

  // Parse GLSL to AST (abstract syntax tree)
  if (!shader.parse(&DefaultTBuiltInResource, 450, false, EShMsgDefault)) {
    std::cerr << "GLSL parsing failed: " << shader.getInfoLog() << std::endl;
    return {};
  }

  // Create a program to hold the compiled shader
  glslang::TProgram program;
  program.addShader(&shader);

  // Link the program
  if (!program.link(EShMsgDefault)) {
    std::cerr << "Linking failed: " << program.getInfoLog() << std::endl;
    return {};
  }

  // Convert AST to SPIR-V
  std::vector<uint32_t> spirv;
  glslang::SpvOptions   spvOptions;
  glslang::GlslangToSpv(*program.getIntermediate(shaderType), spirv, &spvOptions);

  return spirv;
}
