#include <png.h>
#include <EGL/egl.h>
#include <GL/gl.h>
#include <GL/glext.h>
#include <GLES3/gl32.h>

#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

const uint32_t kWidth = 255;
const uint32_t kHeight = 255;
const uint32_t kLayers = 10;
const std::string kVertexShaderPath = "geometry.vert";
const std::string kGeometryShaderPath = "geometry.geom";
const std::string kFragmentShaderPath = "geometry.frag";

const std::string kBlitVertexPath = "blit.vert";
const std::string kBlitFragmentPath = "blit.frag";

EGLDisplay display;
EGLContext context;
GLuint program;

#define ERROR(message) \
    std::cerr << message << std::endl; \
    std::exit(EXIT_FAILURE)

void InitializeEGL() {
  display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
  eglInitialize(display, nullptr, nullptr);
  eglBindAPI(EGL_OPENGL_ES_API);

  const EGLint config_attribs[] = {
      EGL_SURFACE_TYPE, EGL_DONT_CARE,
      EGL_RENDERABLE_TYPE, EGL_OPENGL_ES2_BIT,
      EGL_NONE };
  EGLConfig egl_config;
  EGLint num_configs;
  if (!eglChooseConfig(display, config_attribs, &egl_config, 1, &num_configs)) {
    ERROR("Unable to choose EGL config");
  }

  const EGLint context_attribs[] = { EGL_CONTEXT_CLIENT_VERSION, 2, EGL_NONE };
  context = eglCreateContext(display, egl_config, EGL_NO_CONTEXT, context_attribs);
  if (context == EGL_NO_CONTEXT) {
    ERROR("Failed to create GLES context");
  }

  if (!eglMakeCurrent(display, EGL_NO_SURFACE, EGL_NO_SURFACE, context)) {
    ERROR("Failed to make GLES context current");
  }
  std::cout << glGetString(GL_VERSION) << std::endl;
}

void SetUpLayeredFramebuffer(GLuint layered_texture) {
  glBindTexture(GL_TEXTURE_2D_ARRAY, layered_texture);
  glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
  glTexParameteri(GL_TEXTURE_2D_ARRAY, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
  glTexImage3D(GL_TEXTURE_2D_ARRAY, 0, GL_RGBA8, kWidth, kHeight, kLayers, 0,
      GL_RGBA, GL_UNSIGNED_BYTE, nullptr);
  glBindTexture(GL_TEXTURE_2D_ARRAY, 0);

  GLuint framebuffer;
  glGenFramebuffers(1, &framebuffer);
  glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);
  glFramebufferTexture(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, layered_texture, 0);
}

void SetUpBasicFramebuffer() {
  GLuint color_buffer;
  glGenRenderbuffers(1, &color_buffer);
  glBindRenderbuffer(GL_RENDERBUFFER, color_buffer);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_RGBA8, kWidth, kHeight);

  GLuint framebuffer;
  glGenFramebuffers(1, &framebuffer);
  glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_RENDERBUFFER, color_buffer);
}

GLuint CompileShader(GLenum type, std::string path) {
  std::ifstream fstream(path);
  if (!fstream) {
    ERROR("Unable to read shader: " << path);
  }
  std::stringstream buffer;
  buffer << fstream.rdbuf();
  // TODO(brkho): Read the shader source in more intelligently.
  std::string shader_source_tmp = buffer.str();
  const char* shader_source = shader_source_tmp.c_str();

  GLuint shader = glCreateShader(type);
  if (!shader) {
    ERROR("Failed to create shader of type: " << type);
  }
  glShaderSource(shader, 1, &shader_source, nullptr);
  glCompileShader(shader);
  GLint status;
  glGetShaderiv(shader, GL_COMPILE_STATUS, &status);
  if (!status) {
    GLint log_length;
    glGetShaderiv(shader, GL_INFO_LOG_LENGTH, &log_length);
    char shader_log[log_length];
    glGetShaderInfoLog(shader, log_length, nullptr, shader_log);
    ERROR("Failed to compile shader: " << shader_log);
  }
  return shader;
}

GLuint CreateProgram() {
  GLuint vertex_shader = CompileShader(GL_VERTEX_SHADER, kVertexShaderPath);
  GLuint geometry_shader = CompileShader(GL_GEOMETRY_SHADER, kGeometryShaderPath);
  GLuint fragment_shader = CompileShader(GL_FRAGMENT_SHADER, kFragmentShaderPath);

  GLuint program = glCreateProgram();
  if (!program) {
    ERROR("Failed to create program");
  }
  glAttachShader(program, vertex_shader);
  glAttachShader(program, geometry_shader);
  glAttachShader(program, fragment_shader);
  glLinkProgram(program);
  GLint status;
  glGetProgramiv(program, GL_LINK_STATUS, &status);
  if (!status) {
    GLint log_length = 0;
    glGetProgramiv(program, GL_INFO_LOG_LENGTH, &log_length);
    char program_log[log_length];
    glGetProgramInfoLog(program, log_length, nullptr, program_log);
    ERROR("Failed to link program: " << program_log);
  }
  glDetachShader(program, vertex_shader);
  glDetachShader(program, geometry_shader);
  glDetachShader(program, fragment_shader);
  glDeleteShader(vertex_shader);
  glDeleteShader(geometry_shader);
  glDeleteShader(fragment_shader);

  return program;
}

GLuint CreateBlitProgram() {
  GLuint vertex_shader = CompileShader(GL_VERTEX_SHADER, kBlitVertexPath);
  GLuint fragment_shader = CompileShader(GL_FRAGMENT_SHADER, kBlitFragmentPath);

  GLuint program = glCreateProgram();
  if (!program) {
    ERROR("Failed to create program");
  }
  glAttachShader(program, vertex_shader);
  glAttachShader(program, fragment_shader);
  glLinkProgram(program);
  GLint status;
  glGetProgramiv(program, GL_LINK_STATUS, &status);
  if (!status) {
    GLint log_length = 0;
    glGetProgramiv(program, GL_INFO_LOG_LENGTH, &log_length);
    char program_log[log_length];
    glGetProgramInfoLog(program, log_length, nullptr, program_log);
    ERROR("Failed to link program: " << program_log);
  }
  glDetachShader(program, vertex_shader);
  glDetachShader(program, fragment_shader);
  glDeleteShader(vertex_shader);
  glDeleteShader(fragment_shader);

  return program;
}

void Draw(GLuint program) {
  glViewport(0, 0, kWidth, kHeight);
  glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
  glClear(GL_COLOR_BUFFER_BIT);

  glUseProgram(program);
  glDrawArrays(GL_POINTS, 0, 1);
  glFinish();
}

void Blit(GLuint program, GLuint layered_texture) {
  glViewport(0, 0, kWidth, kHeight);
  glClearColor(0.0f, 0.0f, 0.0f, 1.0f);
  glClear(GL_COLOR_BUFFER_BIT);

  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D_ARRAY, layered_texture);
  glUseProgram(program);
  glDrawArrays(GL_TRIANGLES, 0, 6);
  glFinish();
}

GLubyte* ReadFramebuffer() {
  // TODO(brkho): Maybe use a std::array instead?
  GLubyte* pixels = static_cast<GLubyte*>(malloc(kWidth * kHeight * 4));
  glReadPixels(0, 0, kWidth, kHeight, GL_RGBA, GL_UNSIGNED_BYTE, pixels);
  return pixels;
}

void WriteImage(GLubyte* pixels) {
  FILE *f = fopen("geometry.png", "wb");

  png_image image_info = {};
  image_info.version = PNG_IMAGE_VERSION;
  image_info.width = kWidth;
  image_info.height = kHeight;
  image_info.format = PNG_FORMAT_RGBA;
  if (png_image_write_to_stdio(&image_info, f, 0, pixels, kWidth * 4, nullptr) == 0) {
    ERROR("Error writing PNG: " << image_info.message);
  }

  fclose(f);
}

int main() {
  InitializeEGL();
  GLuint layered_texture;
  glGenTextures(1, &layered_texture);
  SetUpLayeredFramebuffer(layered_texture);
  GLuint program = CreateProgram();
  Draw(program);

  SetUpBasicFramebuffer();
  GLuint blit_program = CreateBlitProgram();
  Blit(blit_program, layered_texture);

  GLubyte* pixels = ReadFramebuffer();
  WriteImage(pixels);

  eglMakeCurrent(display, EGL_NO_SURFACE,EGL_NO_SURFACE, EGL_NO_CONTEXT);
  eglDestroyContext(display, context);
  eglTerminate(display);
  return EXIT_SUCCESS;
}