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

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

const uint32_t kWidth = 255;
const uint32_t kHeight = 255;
const std::string kVertexShaderPath = "occlusion.vert";
const std::string kFragmentShaderPath = "occlusion.frag";

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

void InitializeEGL() {
  EGLDisplay 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 };
  EGLContext 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");
  }
}

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

  GLuint depth_buffer;
  glGenRenderbuffers(1, &depth_buffer);
  glBindRenderbuffer(GL_RENDERBUFFER, depth_buffer);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT16, kWidth, kHeight);

  GLuint framebuffer;
  glGenFramebuffers(1, &framebuffer);
  glBindFramebuffer(GL_FRAMEBUFFER, framebuffer);
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_RENDERBUFFER, color_buffer);
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, depth_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 fragment_shader = CompileShader(GL_FRAGMENT_SHADER, kFragmentShaderPath);

  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;
}

uint32_t Draw(GLuint program) {
  glViewport(0, 0, kWidth, kHeight);
  glClearColor(0.0f, 0.0f, 1.0f, 1.0f);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  glEnable(GL_DEPTH_TEST);

  glUseProgram(program);
  GLuint query_object;
  glGenQueries(1, &query_object);
  glBeginQuery(GL_SAMPLES_PASSED, query_object);
  if (glGetError()) {
    ERROR("GLES version does not support GL_SAMPLES_PASSED");
  }
  glDrawArrays(GL_TRIANGLES, 0, 6);
  glEndQuery(GL_SAMPLES_PASSED);

  uint32_t shaded_fragments = 0;
  glGetQueryObjectuiv(query_object, GL_QUERY_RESULT, &shaded_fragments);

  glFinish();
  return shaded_fragments;
}

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("occlusion.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();
  SetUpFramebuffer();
  GLuint program = CreateProgram();
  
  uint32_t shaded_fragments = Draw(program);
  std::cout << "[Occlusion query] shaded fragments: " << shaded_fragments << std::endl;

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

  return EXIT_SUCCESS;
}