#include <stdbool.h>
#include <stdlib.h>

#include <openssl/rand.h>
#include <GL/glew.h>
#include <SDL2/SDL.h>

#define UNUSED(a) ((void) (a))

// string utility macros
#define SS(v) #v
#define S(v) SS(v)

// number of time records to keep (must be power of two - 1)
#define MAX_TIMES 0xFF

// custom SDL_USEREVENT code for timer
#define EVENT_CODE_TIMER 0

// use compute shader to update particles?
#define USE_COMPUTE_SHADER false

// number of particles
#define NUM_PARTICLES 1000

// number of compute shader workgroups
#define WORKGROUP_SIZE 1

typedef struct {
  const GLuint type;
  const char *src;
} shader_t;

typedef struct {
  SDL_Window *win;

  // times ring buffer
  struct {
    uint32_t draw_start,
             draw_end;
  } times[MAX_TIMES];
  uint32_t times_ofs;

  bool done,
       fullscreen;
} context_t;

static const float
verts[] = {
  -1,  1, 0,
   1,  1, 0,
   0.0, -1, 0,
};

static const char
cs_src[] =
  "#version 320 es\n"
  "\n"
  "layout(local_size_x = " S(WORKGROUP_SIZE) ") in;\n"
  "\n"
  "layout (std430, binding = 1) buffer particles {\n"
  "  mediump vec4 positions[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 velocities[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 colors[" S(NUM_PARTICLES) "];\n"
  "};\n"
  "\n"
  // "// layout (location = 0) in mediump float delta;\n"
  // "\n"
  "void main() {\n"
  // "  return;\n"
  "  mediump uint i = gl_GlobalInvocationID.x;\n"
  "  positions[i].xy += 0.016 * velocities[i].xy;\n"
  "}\n";

static const char
vs_src[] =
  "#version 320 es\n"
  "\n"
  "layout (std430, binding = 1) buffer particles {\n"
  "  mediump vec4 positions[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 velocities[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 colors[" S(NUM_PARTICLES) "];\n"
  "};\n"
  "\n"
  "layout (location = 0) in mediump vec3 pos;\n"
  "out mediump flat int instance_id;\n"
  "\n"
  "void main() {\n"
  "  gl_Position = positions[gl_InstanceID] + 0.025 * vec4(pos, 1);\n"
  "  instance_id = gl_InstanceID;\n"
  "}\n";

static const char
fs_src[] =
  "#version 320 es\n"
  "\n"
  "layout (std430, binding = 1) buffer particles {\n"
  "  mediump vec4 positions[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 velocities[" S(NUM_PARTICLES) "];\n"
  "  mediump vec4 colors[" S(NUM_PARTICLES) "];\n"
  "};\n"
  "\n"
  "in mediump flat int instance_id;\n"
  "out mediump vec4 color;\n"
  "\n"
  "uniform mediump float time;\n"
  "\n"
  "void main() {\n"
  "  mediump float alpha = 0.5f + 0.5f * (0.5f + 0.5 * sin(time * 3.141));\n"
  "  color = vec4(colors[instance_id].rgb, 1) * alpha;\n"
  "}\n";

static const shader_t
COMPUTE_SHADERS[] = {
  { GL_COMPUTE_SHADER, cs_src },
};

static const shader_t
RENDER_SHADERS[] = {
  { GL_VERTEX_SHADER,   vs_src },
  { GL_FRAGMENT_SHADER, fs_src },
};

static GLuint
compile_shader(
  const GLenum type,
  const char *src
) {
  // create shader
  GLuint r = glCreateShader(type);

  // compile shader
  glShaderSource(r, 1, &src, NULL);
  glCompileShader(r);

  // check for error
  GLint ok;
  glGetShaderiv(r, GL_COMPILE_STATUS, &ok);
  if (!ok) {
    GLchar log[1024];
    glGetShaderInfoLog(r, sizeof(log), NULL, log);
    SDL_Log("glCompileShader() failed: %s", log);
    exit(EXIT_FAILURE);
  }

  // return result
  return r;
}

static GLuint
link_program(
  const size_t num_shaders,
  const shader_t *shaders
) {
  // create program
  GLuint r = glCreateProgram();

  // compile shaders
  // (FIXME: check num_shaders < 128)
  GLuint ids[128];
  if (num_shaders >= 128) {
  }
  for (size_t i = 0; i < num_shaders; i++) {
    // compile shader
    ids[i] = compile_shader(shaders[i].type, shaders[i].src);

    // attach shader
    glAttachShader(r, ids[i]);
  }

  // link program
  glLinkProgram(r);

  // check link status
  GLint ok;
  glGetProgramiv(r, GL_LINK_STATUS, &ok);
  if (!ok) {
    GLchar log[1024];
    glGetProgramInfoLog(r, sizeof(log), NULL, log);
    SDL_Log("glLinkProgram() failed: %s", log);
    exit(EXIT_FAILURE);
  }

  // delete shaders
  for (size_t i = 0; i < num_shaders; i++) {
    glDeleteShader(ids[i]);
  }

  // return result
  return r;
}

static float
color_map(uint8_t bits) {
  switch (bits & 0x3) {
  case 0x3:
    return 1.0;
  case 0x2:
    return 0.75;
  case 0x1:
    return 0.5;
  default:
    return 0.25;
  }
}

static GLuint
init_particles(void) {
  // create shader storage buffer
  GLuint ssbo;
  glGenBuffers(1, &ssbo);

  // generate seed data
  uint8_t seeds[5 * NUM_PARTICLES];
  if (!RAND_bytes(seeds, 5 * NUM_PARTICLES)) {
    SDL_Log("RAND_bytes() failed");
    exit(EXIT_FAILURE);
  }

  // populate data
  float data[3 * 4 * NUM_PARTICLES];
  #pragma omp parallel for
  for (int i = 0; i < NUM_PARTICLES; i++) {
    float x = 2.0 * seeds[5 * i + 0] / 0xFF - 1.0,
          y = 2.0 * seeds[5 * i + 1] / 0xFF - 1.0;

    // SDL_Log("%f,%f", x, y);

    // populate position
    data[4 * i + 0] = x;
    data[4 * i + 1] = y;
    data[4 * i + 2] = 0.0;
    data[4 * i + 3] = 1.0;

    // populate random velocity
    data[4 * (NUM_PARTICLES + i) + 0] = 2.0 * seeds[5 * i + 2] / 0xFF - 1.0;
    data[4 * (NUM_PARTICLES + i) + 1] = 2.0 * seeds[5 * i + 3] / 0xFF - 1.0;

    // populate colors
    uint8_t cs = seeds[5 * i + 4];
    data[4 * (2 * NUM_PARTICLES + i) + 0] = color_map(cs & 0x3);
    data[4 * (2 * NUM_PARTICLES + i) + 1] = color_map((cs >> 2) & 0x3);
    data[4 * (2 * NUM_PARTICLES + i) + 2] = color_map((cs >> 4) & 0x3);
    data[4 * (2 * NUM_PARTICLES + i) + 3] = 1.0;
  }

  // populate buffer
  glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
  glBufferData(GL_SHADER_STORAGE_BUFFER, sizeof(data), data, GL_DYNAMIC_DRAW);
  glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);

  // return buffer
  return ssbo;
}

static void
update_particles(
  const GLuint ssbo,
  const GLuint compute_prog
) {
  if (USE_COMPUTE_SHADER) {
    UNUSED(ssbo);

    glUseProgram(compute_prog);
    glDispatchCompute(NUM_PARTICLES / WORKGROUP_SIZE, 1, 1);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
  } else {
    UNUSED(compute_prog);

    // bind shader storage buffer
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);

    // map buffer
    float *data = glMapBufferRange(
      GL_SHADER_STORAGE_BUFFER,
      0,
      3 * NUM_PARTICLES * 4 * sizeof(float),
      GL_MAP_READ_BIT | GL_MAP_WRITE_BIT
    );

    // update particle positions
    #pragma omp parallel for
    for (int i = 0; i < NUM_PARTICLES; i++) {
      // update x coordinate
      data[4 * i + 0] += 0.016 * data[4 * (NUM_PARTICLES + i) + 0];
      if (data[4 * i + 0] < -1.1) {
        data[4 * i + 0] += 2.2;
      } else if (data[4 * i + 0] > 1.1) {
        data[4 * i + 0] -= 2.2;
      }

      // update y coordinate
      data[4 * i + 1] += 0.016 * data[4 * (NUM_PARTICLES + i) + 1];
      if (data[4 * i + 1] < -1.1) {
        data[4 * i + 1] += 2.2;
      } else if (data[4 * i + 1] > 1.1) {
        data[4 * i + 1] -= 2.2;
      }
    }

    // unmap and unbind data buffer
    glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
  }
}

static void
update_viewport(
  SDL_Window *win
) {
  int w, h;

  // get window drawable size
  SDL_GL_GetDrawableSize(win, &w, &h);

  // set viewport size
  glViewport(0, 0, w, h);
}

static void
ctx_init(context_t * const ctx) {
  // init sdl
  if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_TIMER) < 0) {
    SDL_Log("SDL_Init() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // create window
  ctx->win = SDL_CreateWindow(
    "Compute Test",
    SDL_WINDOWPOS_UNDEFINED,
    SDL_WINDOWPOS_UNDEFINED,
    640,
    480,
    SDL_WINDOW_OPENGL | SDL_WINDOW_RESIZABLE
  );

  // check for error
  if (!ctx->win) {
    SDL_Log("SDL_CreateWindow() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // set gl context majory version hint
  if (SDL_GL_SetAttribute(
    SDL_GL_CONTEXT_PROFILE_MASK,
    SDL_GL_CONTEXT_PROFILE_ES
  ) < 0) {
    SDL_Log("SDL_GL_SetAttribute() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // set gl context majory version hint
  if (SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 3) < 0) {
    SDL_Log("SDL_GL_SetAttribute() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // set gl context minor version hint
  if (SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 2) < 0) {
    SDL_Log("SDL_GL_SetAttribute() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // create GL context
  SDL_GLContext context = SDL_GL_CreateContext(ctx->win);
  if (!context) {
    SDL_Log("SDL_GL_CreateContext() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // init glew
  // (do this immediately after GL context init)
  {
    GLenum err = glewInit();
    if (err != GLEW_OK) {
      SDL_Log("glewInit() failed: %s", glewGetErrorString(err));
      exit(EXIT_FAILURE);
    }
  }

  // set viewport size
  // (after GLEW init)
  update_viewport(ctx->win);

  // set swap interval
  // (after context init)
  if (SDL_GL_SetSwapInterval(1) < 0) {
    SDL_Log("SDL_GL_SetSwapInterval() failed: %s", SDL_GetError());
    exit(EXIT_FAILURE);
  }

  // init times ring buffer
  // memset(ctx->times, 0, sizeof(ctx->times));
  ctx->times_ofs = 0;

  // init flags
  ctx->done = false;
  ctx->fullscreen = false;
}

static void
handle_events(
  context_t * const ctx
) {
  SDL_Event ev;

  while (SDL_PollEvent(&ev)) {
    switch (ev.type) {
    case SDL_QUIT:
      ctx->done = true;
      break;
    case SDL_KEYDOWN:
      switch (ev.key.keysym.sym) {
      case SDLK_q:
      case SDLK_ESCAPE:
        ctx->done = true;

        break;
      case SDLK_F11:
        // toggle fullscreen
        ctx->fullscreen = !ctx->fullscreen;

        if (SDL_SetWindowFullscreen(
          ctx->win,
          ctx->fullscreen ? SDL_WINDOW_FULLSCREEN_DESKTOP : 0
        ) < 0) {
          // log warning
          SDL_Log("SDL_SetWindowFullscreen() failed: %s", SDL_GetError());
          ctx->fullscreen = !ctx->fullscreen;
        }

        break;
      }

      break;
    case SDL_WINDOWEVENT:
      switch (ev.window.event) {
      case SDL_WINDOWEVENT_SIZE_CHANGED:
        {
          // map event to window
          SDL_Window *ev_win = SDL_GetWindowFromID(ev.window.windowID);

          if (ev_win == ctx->win) {
            // update viewport
            update_viewport(ctx->win);
          } else {
            SDL_Log("ignoring SIZE_CHANGED event from unknown window");
          }
        }

        break;
      }

      break;
    case SDL_USEREVENT:
      if (ev.user.code == EVENT_CODE_TIMER) {
        uint32_t sum = 0;

        // #pragma omp parallel for reduction(+:sum)
        for (int i = 0; i < MAX_TIMES; i++) {
          sum += ctx->times[i].draw_end - ctx->times[i].draw_start;
        }

        // print fps
        SDL_Log("fps: %f", sum * 1.0 / MAX_TIMES);
      }

      break;
    }
  }
}

static Uint32
timer_cb(
  Uint32 interval,
  void *arg
) {
  SDL_Event ev;
  UNUSED(arg);

  // build synthetic event
  ev.type = SDL_USEREVENT;
  ev.user.type = SDL_USEREVENT;
  ev.user.code = EVENT_CODE_TIMER;
  ev.user.data1 = NULL;
  ev.user.data2 = NULL;

  // add event to queue
  SDL_PushEvent(&ev);

  // return interval
  return interval;
}

int main(int argc, char *argv[]) {
  UNUSED(argc);
  UNUSED(argv);

  // init context
  context_t ctx;
  ctx_init(&ctx);

  // init timer (5 seconds)
  SDL_TimerID timer_id = SDL_AddTimer(5000, timer_cb, NULL);

  // generate vertex array
  GLuint vao;
  glGenVertexArrays(1, &vao);
  glBindVertexArray(vao);

  // generate shader storage buffer
  GLuint ssbo = init_particles();
  glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssbo);

  // generate vertex buffer
  GLuint vbo;
  glGenBuffers(1, &vbo);
  glBindBuffer(GL_ARRAY_BUFFER, vbo);
  glBufferData(GL_ARRAY_BUFFER, sizeof(verts), verts, GL_STATIC_DRAW);

  glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 3 * sizeof(GLfloat), (GLvoid*) 0);
  glEnableVertexAttribArray(0);

  // unbind vao
  glBindVertexArray(0);

  // link compute program
  GLuint compute_prog = link_program(1, COMPUTE_SHADERS);
  // GLint u_delta = glGetUniformLocation(compute_prog, "delta");

  // link render program, get uniforms
  GLuint render_prog = link_program(2, RENDER_SHADERS);
  GLint u_time = glGetUniformLocation(render_prog, "time");

  // main loop
  while (!ctx.done) {
    // get start time
    uint32_t now = SDL_GetTicks();

    // save start time
    ctx.times[ctx.times_ofs].draw_start = now;

    // update particles
    update_particles(ssbo, compute_prog);

    // clear screen
    glClearColor(0, 0, 0, 1);
    glClear(GL_COLOR_BUFFER_BIT);

    // use render program, set uniform
    glUseProgram(render_prog);
    glUniform1f(u_time, now / 1000.0);

    // draw
    glBindVertexArray(vao);
    glDrawArraysInstanced(GL_TRIANGLES, 0, 3, NUM_PARTICLES);
    glBindVertexArray(0);

    // swap buffer
    SDL_GL_SwapWindow(ctx.win);

    // save end time
    ctx.times[ctx.times_ofs].draw_end = SDL_GetTicks();
    ctx.times_ofs = (ctx.times_ofs + 1) & MAX_TIMES;

    // handle events
    handle_events(&ctx);
  }

  // remove timer
  if (SDL_RemoveTimer(timer_id) == SDL_FALSE) {
    SDL_Log("SDL_RemoveTimer() failed");
  }

  // delete GL context
  SDL_GLContext context = SDL_GL_GetCurrentContext();
  if (context) {
    SDL_GL_DeleteContext(context);
  }

  // destroy window
  SDL_DestroyWindow(ctx.win);

  // return success
  return EXIT_SUCCESS;
}