#include <vector>
#include <cmath>
#include <sstream>
#include <iostream>
#include <glad/glad.h>
#include <GLFW/glfw3.h>

GLFWAPI GLFWwindow* createWindow(int width, int height, const char* title) {
  GLFWwindow* window = nullptr;
  glfwSetErrorCallback([](int /*error_code*/, const char* description) {
    std::cerr << description << std::endl;
    std::exit(EXIT_FAILURE);
    });
  glfwInit();

  glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
  glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 6);
  glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
  glfwWindowHint(GLFW_OPENGL_DEBUG_CONTEXT, true);
  window = glfwCreateWindow(width, height, title, NULL, NULL);

  if (!window) {
    std::cerr << "Failed to create GLFW window" << std::endl;
    glfwTerminate();
    std::exit(EXIT_FAILURE);
  }

  // Make the OpenGL context for this window be the currently associated context for this thread.
  glfwMakeContextCurrent(window);

  // Load the OpenGL API function pointers.
  if (!gladLoadGLLoader((GLADloadproc)glfwGetProcAddress)) {
    std::cerr << "Failed to initialize GLAD" << std::endl;
    glfwDestroyWindow(window);
    glfwTerminate();
    std::exit(EXIT_FAILURE);
  }

  return window;
}

void checkShaderProgram(GLuint shader_program) {
  GLint status = GL_FALSE;
  glGetProgramiv(shader_program, GL_LINK_STATUS, &status);
  if (status == GL_FALSE) {
    GLchar info_log[4096];
    glGetProgramInfoLog(shader_program, sizeof(info_log), NULL, info_log);
    std::cerr << info_log << std::endl;
    std::exit(EXIT_FAILURE);
  }
};

GLuint createShaderProgram(GLenum shader_type, const std::vector<const GLchar*>& shader_sources) {
  GLuint shader_program = glCreateShaderProgramv(shader_type, (GLsizei)shader_sources.size(), shader_sources.data());
  checkShaderProgram(shader_program);
  return shader_program;
}

int main()
{
  GLFWwindow* window = createWindow(1024, 1024, "Lesson 3: Textures");

  const size_t num_instances = 16, num_draws = num_instances * num_instances;

  GLuint sampler = 0;
  GLuint textures[num_draws];
  GLuint64 texture_handles[num_draws];
  GLuint texture_handles_buffer = 0;
  {
    glCreateSamplers(1, &sampler);
    glSamplerParameteri(sampler, GL_TEXTURE_MAG_FILTER, GL_NEAREST);

    glCreateTextures(GL_TEXTURE_2D_ARRAY, num_draws, textures);
    const size_t texture_width = 16, texture_height = 16, texture_depth = num_instances;
    struct Texel {
      GLubyte r;
      GLubyte g;
      GLubyte b;
      GLubyte a;
    };

    // Convert linear sRGB to gamma-corrected sRGB.
    auto linearToSrgb = [](float c) -> GLubyte {
      if (c <= 0.0031308f)
        c *= 12.92f;
      else
        c = 1.055f * std::pow(c, 1.f / 2.4f) - 0.055f;
      c = c > 1.f ? 1.f : c < 0.f ? 0.f : c;
      return (GLubyte)(c * 255.f);
    };

    auto hsvToRgb = [&](float h, float s, float v) -> Texel {
      float r, g, b;
      int i = int(h * 6);
      float f = h * 6 - i;
      float p = v * (1 - s);
      float q = v * (1 - f * s);
      float t = v * (1 - (1 - f) * s);
      switch (i % 6) {
        case 0: r = v, g = t, b = p; break;
        case 1: r = q, g = v, b = p; break;
        case 2: r = p, g = v, b = t; break;
        case 3: r = p, g = q, b = v; break;
        case 4: r = t, g = p, b = v; break;
        case 5: r = v, g = p, b = q; break;
      }
      return Texel{
        linearToSrgb(r),
        linearToSrgb(g),
        linearToSrgb(b),
        255
      };
    };

    // For each texture, draw a different colored cones as discs stacked for each array of 2D slices.
    for (size_t i = 0; i < num_draws; i++) {
      float draw_ratio = float(i) / num_draws;
      glTextureStorage3D(textures[i], 1, GL_SRGB8_ALPHA8, texture_width, texture_height, texture_depth);

      for (size_t z = 0; z < texture_depth; z++) {
        Texel image[texture_height][texture_width];
        float depth_ratio = float(z) / texture_depth;
        float radius = 1.f - depth_ratio;
        Texel color = hsvToRgb(draw_ratio, 1.0f, depth_ratio);

        for (size_t y = 0; y < texture_height; y++) {
          float yf = (float(y) + 0.5f) / texture_height * 2.f - 1.f;

          for (size_t x = 0; x < texture_width; x++) {
            float xf = (float(x) + 0.5f) / texture_width * 2.f - 1.f;
            image[y][x] = color;
            if ((xf * xf) + (yf * yf) >= (radius * radius))
              image[y][x].a = 0;
          }
        }
        GLint level = 0, xoffset = 0, yoffset = 0, zoffset = (GLint)z;
        GLsizei width = texture_width, height = texture_height, depth = 1;
        glTextureSubImage3D(textures[i], level, xoffset, yoffset, zoffset, width, height, depth, GL_RGBA, GL_UNSIGNED_BYTE, image);
      }
      glGenerateTextureMipmap(textures[i]);
      // Get the GPU address of the texture and make it resident so shaders can access it.
      texture_handles[i] = glGetTextureSamplerHandleARB(textures[i], sampler);
      glMakeTextureHandleResidentARB(texture_handles[i]);
    }

    glCreateBuffers(1, &texture_handles_buffer);
    glNamedBufferStorage(texture_handles_buffer, sizeof(texture_handles), texture_handles, 0);
  }

  struct Vertex {
    float x;
    float y;
    float u;
    float v;
  };

  struct Triangle {
    unsigned short v0;
    unsigned short v1;
    unsigned short v2;
  };

  float quad_width = 1.0f / num_instances;
  Vertex quad_verts[4] {
    {-quad_width, -quad_width, 0.0f, 0.f},
    {+quad_width, -quad_width, 1.0f, 0.f},
    {+quad_width, +quad_width, 1.0f, 1.f},
    {-quad_width, +quad_width, 0.0f, 1.f},
  };
  Triangle quad_indices[2]{
    {0, 1, 2},
    {0, 2, 3}
  };

  struct Mesh {
    unsigned int instance_count = 0;
    unsigned int base_index = 0;
    unsigned int index_count = 0;
    unsigned int base_vertex = 0;
    unsigned int vertex_count = 0;
  };
  std::vector<Mesh> meshes;


  GLuint mesh_buffer = 0;
  GLint verts_buffer_offset = 0;
  glCreateBuffers(1, &mesh_buffer);
  {
    glNamedBufferStorage(mesh_buffer, sizeof(quad_indices) + sizeof(quad_verts), nullptr, GL_MAP_WRITE_BIT);
    {
      GLubyte* mapped = (GLubyte*)glMapNamedBuffer(mesh_buffer, GL_WRITE_ONLY);
      memcpy(mapped, quad_indices, sizeof(quad_indices));
      verts_buffer_offset = (GLint)sizeof(quad_indices);
      memcpy(mapped + verts_buffer_offset, quad_verts, sizeof(quad_verts));
      glUnmapNamedBuffer(mesh_buffer);
    }

    Mesh mesh{ num_instances, 0, (unsigned int)std::size(quad_indices) * 3, 0, (unsigned int)std::size(quad_verts) };
    meshes = std::vector<Mesh>(num_draws, mesh);
  }

  GLuint vertex_array_object = 0;
  glCreateVertexArrays(1, &vertex_array_object);

  const GLuint mesh_buffer_binding = 0, position_attrib = 0, texcoord_attrib = 1;
  glVertexArrayAttribBinding(vertex_array_object, position_attrib, mesh_buffer_binding);
  glVertexArrayAttribFormat( vertex_array_object, position_attrib, 2, GL_FLOAT, GL_FALSE, offsetof(Vertex, x));
  glEnableVertexArrayAttrib( vertex_array_object, position_attrib);

  glVertexArrayAttribBinding(vertex_array_object, texcoord_attrib, mesh_buffer_binding);
  glVertexArrayAttribFormat( vertex_array_object, texcoord_attrib, 2, GL_FLOAT, GL_FALSE, offsetof(Vertex, u));
  glEnableVertexArrayAttrib( vertex_array_object, texcoord_attrib);


  struct Vec2 {
    GLfloat x;
    GLfloat y;
  };
  GLuint animation_buffer = 0;
  glCreateBuffers(1, &animation_buffer);
  glNamedBufferStorage(animation_buffer, sizeof(Vec2[num_draws][num_instances]), nullptr, GL_DYNAMIC_STORAGE_BIT);


  struct ComputeShader {
    GLuint program;
    GLint frame_uniform;
  } compute_shader;

  struct VertexShader {
    GLuint program;
    GLint rotation_uniform;
  } vertex_shader;

  struct FragmentShader {
    GLuint program;
  } fragment_shader;

  {
    // The version must be the first statement in the shader.
    const GLchar* shader_version_source = R"(
      #version 460 core // OpenGL 4.6 Core Profile
    )";

    // Any extensions must be enabled immediately after the version declaration.
    const GLchar* vertex_shader_extensions = R"(
        #define IN_OUT out
    )";
    const GLchar* fragment_shader_extensions = R"(
      #extension GL_ARB_bindless_texture : require
      #define IN_OUT in
    )";

    GLchar compute_shader_source[4096];
    snprintf(compute_shader_source, std::size(compute_shader_source),  R"(
        // Compute shader to generate a rotation matrix.
        layout(local_size_x = %zd, local_size_y = 1, local_size_z = 1) in;

        uniform uint frame;

        layout(binding = 1, std430) writeonly buffer ssbo { vec2 animations[]; };

        void main() {
          uint num_meshes = gl_NumWorkGroups.x, mesh = gl_WorkGroupID.x;
          uint num_instances = gl_WorkGroupSize.x, instance = gl_LocalInvocationID.x;

          uint num_columns = num_instances, num_rows = num_instances;
          uint row = uint(mesh / num_columns), column = uint(mod(mesh, num_columns));
          float column_offset = bool(row & 1) ? 0.75f : 0.25f;
          float mesh_ratio_x = (column + column_offset) / num_columns;
          float mesh_ratio_y = (row + 0.5) / num_rows;

          float instance_ratio = float(instance) / num_instances;
          float stack_offset = instance_ratio / num_rows * 4;

          vec2 offset = vec2(mesh_ratio_x * 2 - 1, mesh_ratio_y * -2 + 1  + stack_offset);
          animations[mesh * num_instances + instance] = offset;
        }
      )", num_instances);

    GLchar vertex_shader_source[4096];
    snprintf(vertex_shader_source, std::size(vertex_shader_source),  R"(
      layout (location = 0) in vec2 position;
      layout (location = 1) in vec2 texcoord;
      layout (binding = 1, std430) readonly buffer ssbo2 { vec2 animations[]; };
      uniform mat2x2 rotation_matrix;

      out gl_PerVertex {
          vec4 gl_Position;
      };

      void main() {
        gl_Position.xy = rotation_matrix * position.xy + animations[gl_BaseInstance + gl_InstanceID];
        //gl_Position.xy += rotation_matrix * vec2(0.0f, 1.0f) * (gl_InstanceID * 0.01f);
        gl_Position.z = 0.0f;
        gl_Position.w = 1.0f;
        vsOutput.texcoord = texcoord.xy;
        int num_instances = %zd;
        draw_id = int(gl_BaseInstance) / num_instances;
        instance_id = int(gl_InstanceID);
      }
    )", num_instances);

    const GLchar* vertex_output_source = R"(
      // Define a common struct that will be used to pass data from the vertex shader to the fragment shader.
      IN_OUT struct VSOutput {
        vec2 texcoord;
      } vsOutput;

      flat IN_OUT int draw_id;
      flat IN_OUT int instance_id;
    )";

    const GLchar* fragment_shader_source = R"(
      layout(binding = 2, std430) readonly buffer ssbo3 {
          sampler2DArray textures[];
      };

      // Just one output value. Automatically goes to the render target.
      out vec4 outColor;

      void main()
      {
        sampler2DArray texture_array = textures[draw_id];
        outColor = texture(texture_array, vec3(vsOutput.texcoord.x, vsOutput.texcoord.y, instance_id));
      }
    )";

    GLuint compute_program = createShaderProgram(GL_COMPUTE_SHADER, { shader_version_source, compute_shader_source });
    GLuint vertex_program = createShaderProgram(GL_VERTEX_SHADER, { shader_version_source, vertex_shader_extensions, vertex_output_source, vertex_shader_source });
    GLuint fragment_program = createShaderProgram(GL_FRAGMENT_SHADER, { shader_version_source, fragment_shader_extensions, vertex_output_source, fragment_shader_source });
    compute_shader = { compute_program, glGetUniformLocation(compute_program, "frame") };
    vertex_shader = { vertex_program, glGetUniformLocation(vertex_program, "rotation_matrix") };
    fragment_shader = { fragment_program };
  }

  GLuint pipeline = 0;
  glGenProgramPipelines(1, &pipeline);
  glUseProgramStages(pipeline, GL_VERTEX_SHADER_BIT, vertex_shader.program);
  glUseProgramStages(pipeline, GL_FRAGMENT_SHADER_BIT, fragment_shader.program);

  constexpr GLuint anim_binding = 1;
  glBindBufferBase(GL_SHADER_STORAGE_BUFFER, anim_binding, animation_buffer);

  constexpr GLuint texture_binding = 2;
  glBindBufferBase(GL_SHADER_STORAGE_BUFFER, texture_binding, texture_handles_buffer);


  // https://registry.khronos.org/OpenGL-Refpages/gl4/html/glMultiDrawArraysIndirect.xhtml
  struct DrawElementsIndirectCommand {
    unsigned int count;
    unsigned int instanceCount;
    unsigned int firstIndex;
    int baseVertex;
    unsigned int baseInstance;
  };
  DrawElementsIndirectCommand cpu_commands_buffer[num_draws] = {};

  GLuint command_buffer = 0;
  glCreateBuffers(1, &command_buffer);
  glNamedBufferStorage(command_buffer, sizeof(cpu_commands_buffer), nullptr, GL_DYNAMIC_STORAGE_BIT);


  for (int frame = 0; !glfwWindowShouldClose(window); frame++) {
    glClearColor(0.0f, 0.0f, 0.5f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);

    glProgramUniform1ui(compute_shader.program, compute_shader.frame_uniform, frame);
    glUseProgram(compute_shader.program);
    glDispatchCompute(num_draws, 1, 1);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glUseProgram(0);

    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    struct Mat2x2 {
      GLfloat m00, m01;
      GLfloat m10, m11;
    };
    float angle = frame * -0.03f;
    Mat2x2 rotation_matrix = {
      std::cos(angle), -std::sin(angle),
      std::sin(angle),  std::cos(angle)
    };
    glProgramUniformMatrix2fv(vertex_shader.program, vertex_shader.rotation_uniform, 1, false, (const float*)&rotation_matrix);
    glBindProgramPipeline(pipeline);

    glBindVertexArray(vertex_array_object);
    const GLsizei bind_count = 1;
    const GLuint bind_buffers[bind_count] = { mesh_buffer };
    const GLintptr bind_offsets[bind_count] = { verts_buffer_offset };
    const GLsizei bind_strides[bind_count] = { sizeof(Vertex) };
    glVertexArrayVertexBuffers(vertex_array_object, mesh_buffer_binding, bind_count, bind_buffers, bind_offsets, bind_strides);
    glVertexArrayElementBuffer(vertex_array_object, mesh_buffer);

    GLuint base_instance = 0;
    GLsizei draw_count = 0;
    for (const Mesh& mesh : meshes) {
      cpu_commands_buffer[draw_count++] = { mesh.index_count, mesh.instance_count, mesh.base_index, (int)mesh.base_vertex, base_instance };
      base_instance += mesh.instance_count;
    }
    glNamedBufferSubData(command_buffer, 0, sizeof(cpu_commands_buffer[0]) * draw_count, cpu_commands_buffer);
    glBindBuffer(GL_DRAW_INDIRECT_BUFFER, command_buffer);

    glMultiDrawElementsIndirect(GL_TRIANGLES, GL_UNSIGNED_SHORT, 0, draw_count, 0);

    glfwSwapBuffers(window);
    glfwPollEvents();
  }

  // We could just let process termination clean everything up for us.
  // But, let's manually clean up our resources just to be explicit.
  glDeleteVertexArrays(1, &vertex_array_object);
  glDeleteProgramPipelines(1, &pipeline);
  glDeleteProgram(compute_shader.program);
  glDeleteProgram(fragment_shader.program);
  glDeleteProgram(vertex_shader.program);
  GLuint buffers_to_delete[] = { command_buffer, animation_buffer, mesh_buffer};
  glDeleteBuffers(std::size(buffers_to_delete), buffers_to_delete);
  glDeleteTextures(num_draws, textures);
  glDeleteSamplers(1, &sampler);

  // Shut down and clean up everything we did with GLFW.
  glfwTerminate();
  // Exit the program.
  return 0;
}