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

GLFWAPI GLFWwindow* createWindow(int width, int height, const char* title);

GLuint createShaderProgram(GLenum shader_type, const std::vector<const GLchar*>& shader_sources);


namespace glsl {
  // Unfortunately, we can't use `alignas` on typedefs for scalars. 
  // So, instead we'll set up shortcut macro for declaring aligned scalar types correctly.
  template < typename Scalar >
  constexpr inline size_t std140_alignof() {
    if constexpr (std::is_same_v<Scalar, float> || std::is_same_v<Scalar, int32_t> || std::is_same_v<Scalar, uint32_t> || std::is_same_v<Scalar, bool>) {
      return 4;
    } else if constexpr (std::is_same_v<Scalar, double> || std::is_same_v<Scalar, int64_t> || std::is_same_v<Scalar, uint64_t>) {
      return 8;
    } else {
      static_assert(sizeof(Scalar) == 0, "Unsupported scalar type for std140 alignment.");
    }
  }

#define STD140(ScalarType)  alignas(glsl::std140_alignof< ScalarType >()) ScalarType

  struct alignas(8) Vec2 {
    STD140(float) x;
    STD140(float) y;
  };
  struct alignas(16) Vec3 {
    STD140(float) x;
    STD140(float) y;
    STD140(float) z;
  };
  struct alignas(16) Vec4 {
    STD140(float) x;
    STD140(float) y;
    STD140(float) z;
    STD140(float) w;
  };

  struct Mat2x2 {
    struct alignas(16) Row { STD140(float) col[2]; };
    Row row[2];
  };
}


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 6: Texture Arrays");

  const int texture_width = 256, texture_height = 256;
  const int num_images_per_array = 32, num_texture_arrays = 16;
  const GLenum texture_format = GL_RGBA8;
  const size_t bytes_per_texel = 4; // RGBA with 8 bits for each channel.
  const size_t image_size = texture_width * texture_height * bytes_per_texel;
  const size_t texture_array_size = image_size * num_images_per_array;

  GLuint texture_arrays[num_texture_arrays];
  {
    glCreateTextures(GL_TEXTURE_2D_ARRAY, num_texture_arrays, texture_arrays);
    for (GLuint texture_array : texture_arrays)
      glTextureStorage3D(texture_array, 1, texture_format, texture_width, texture_height, num_images_per_array);
  }

  const int num_upload_blocks = 4;
  GLuint pixel_unpack_buffer = 0;
  {
    glCreateBuffers(1, &pixel_unpack_buffer);
    const size_t total_buffer_size = texture_array_size * num_upload_blocks;
    glNamedBufferStorage(pixel_unpack_buffer, total_buffer_size, nullptr, GL_MAP_WRITE_BIT | GL_DYNAMIC_STORAGE_BIT);
  }

  {
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, pixel_unpack_buffer);
    for (int array_to_upload = 0; array_to_upload < num_texture_arrays; ++array_to_upload) {
      uintptr_t upload_block_start = (array_to_upload % num_upload_blocks) * texture_array_size;
      {
        GLubyte* mapped = (GLubyte*)glMapNamedBufferRange(pixel_unpack_buffer, upload_block_start, texture_array_size, GL_MAP_WRITE_BIT | GL_MAP_INVALIDATE_RANGE_BIT);
        // Fill the mapped buffer with procedural texture data
        for (int image = 0; image < num_images_per_array; ++image) {
          for (int y = 0; y < texture_height; ++y) {
            for (int x = 0; x < texture_width; ++x) {
              size_t index = image * image_size + (y * texture_width + x) * bytes_per_texel;
              mapped[index + 0] = x % 256; // R
              mapped[index + 1] = ((image ^ array_to_upload)<<4) % 256; // G
              mapped[index + 2] = y % 256; // B
              mapped[index + 3] = 255;     // A
              int cx = x - texture_width/2, cy = y - texture_height / 2;
              if ((cx * cx + cy * cy) < 2048) {
                mapped[index + 0] = 255;
                mapped[index + 2] = 255;
              }
            }
          }
        }
        glUnmapNamedBuffer(pixel_unpack_buffer);
      }
      // Asynchronously upload the texture data from the PBO to the texture array
      glTextureSubImage3D(
        texture_arrays[array_to_upload],
        0, // mip level
        0, 0, 0, // xoffset, yoffset, zoffset
        texture_width, texture_height, num_images_per_array, // width, height, depth
        GL_RGBA,
        GL_UNSIGNED_BYTE,
        (GLvoid*)(upload_block_start)
      );
    }
    glBindBuffer(GL_PIXEL_UNPACK_BUFFER, 0);
  }

  struct Vertex {
    glsl::Vec3 position;
    glsl::Vec2 texcoord;
  };
  Vertex quad_verts[4] = {
      { {-1.0f, -1.0f, 0.0f}, {0.0f, 0.0f} },
      { {+1.0f, -1.0f, 0.0f}, {1.0f, 0.0f} },
      { {-1.0f, +1.0f, 0.0f}, {0.0f, 1.0f} },
      { {+1.0f, +1.0f, 0.0f}, {1.0f, 1.0f} },
  };
  GLuint vertex_buffer = 0;
  glCreateBuffers(1, &vertex_buffer);
  glNamedBufferStorage(vertex_buffer, sizeof(quad_verts), quad_verts, 0);

  // https://www.khronos.org/opengl/wiki/Vertex_Rendering/Rendering_Failure
  // We're using an OpenGL 4.6 Core context. So, we need to set an empty VAO even if we aren't doing anything with it.
  GLuint vertex_array_object = 0;
  glCreateVertexArrays(1, &vertex_array_object);
  glBindVertexArray(vertex_array_object);


  const GLchar* shader_common_source = R"(
    #version 460 core
    // Define a common struct that will be used to pass data from the vertex shader to the fragment shader.
    struct VSOutput {
      vec2 texcoord;
    };
  )";

  const GLchar* vertex_shader_source = R"(
    // Matches the C++ struct above.
    struct Vertex {
      vec3 position;
      vec2 texcoord;
    };

    layout(binding = 0, std430) readonly buffer my_ssbo { Vertex vertices[]; };
    uniform float time;

    out gl_PerVertex {
      vec4 gl_Position;
    };
    out VSOutput vsOutput;

    void main() {
      // Just pass-through the data directly from the SSBO to the outputs.
      vec4 position = vec4(vertices[gl_VertexID].position.xyz, 1.0);
      float angle = time*0.125, sin_t = sin(angle), cos_t = cos(angle);
      position.xy = vec2(position.x * cos_t - position.y*sin_t, position.x*sin_t + position.y*cos_t);
      gl_Position = position;
      vsOutput.texcoord = vertices[gl_VertexID].texcoord;
    }
  )";

  GLchar fragment_shader_source[4096];
  snprintf(fragment_shader_source, std::size(fragment_shader_source), R"(
    // Array of sampler uniforms (bound to texture units 0-15)
    layout(binding = 0) uniform sampler2DArray textures[16];
    uniform float time;

    in VSOutput vsOutput;
    out vec4 fsOutput;

    void main() {
        vec2 uv = vsOutput.texcoord;
        const int num_texture_arrays = %d, num_images_per_array = %d;
        int texture_index = int(uv.x * num_texture_arrays) %% num_texture_arrays;
        int array_element = int(uv.y * num_images_per_array) %% num_images_per_array;
        float angle = time* 0.125, sin_t = sin(angle), cos_t = cos(angle);
        uv.xy = vec2(uv.x * cos_t - uv.y*sin_t, uv.x*sin_t + uv.y*cos_t);
        fsOutput = texture(textures[texture_index], vec3(uv, array_element));
    }
  )", num_texture_arrays, num_images_per_array);

  GLuint vertex_program = createShaderProgram(GL_VERTEX_SHADER, { shader_common_source, vertex_shader_source });
  GLuint fragment_program = createShaderProgram(GL_FRAGMENT_SHADER, { shader_common_source, fragment_shader_source });
  GLuint shader_program_pipeline = 0;
  glGenProgramPipelines(1, &shader_program_pipeline);
  glUseProgramStages(shader_program_pipeline, GL_VERTEX_SHADER_BIT, vertex_program);
  glUseProgramStages(shader_program_pipeline, GL_FRAGMENT_SHADER_BIT, fragment_program);
  GLint vertex_time_location = glGetUniformLocation(vertex_program, "time");
  GLint fragment_time_location = glGetUniformLocation(fragment_program, "time");


  while (!glfwWindowShouldClose(window)) {
    glClearColor(0.0f, 0.0f, 0.5f, 1.0f);
    glClear(GL_COLOR_BUFFER_BIT);

    const GLuint vertex_binding = 0;
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, vertex_binding, vertex_buffer);
    glBindProgramPipeline(shader_program_pipeline);

    float time = static_cast<float>(glfwGetTime());
    glProgramUniform1f(vertex_program, vertex_time_location, time);
    glProgramUniform1f(fragment_program, fragment_time_location, time);

    // Bind all textures to consecutive texture units
    for (int i = 0; i < num_texture_arrays; ++i) {
      glBindTextureUnit(i, texture_arrays[i]);
    }

    glDrawArrays(GL_TRIANGLE_STRIP, 0, 4);

    glfwSwapBuffers(window);
    glfwPollEvents();
  }

  glfwTerminate();
}