#define _GNU_SOURCE
#include <EGL/egl.h>
#include <GL/glew.h>

#include <GL/gl.h>
#include <assert.h>
#include <linux/input-event-codes.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.h>
#include <wayland-client-protocol.h>
#include <wayland-client.h>
#include <wayland-egl-core.h>
#include <wayland-egl.h>
#include <wayland-util.h>
#include <xkbcommon/xkbcommon-keysyms.h>
#include <xkbcommon/xkbcommon-names.h>
#include <xkbcommon/xkbcommon.h>

#include "wlr-layer-shell-unstable-v1.h"
#include "wlr-screencopy-unstable-v1.h"

#include "shaders.h"

static struct wl_display *display;
static EGLDisplay *eglDisplay;
static struct wl_compositor *compositor;
static struct wl_shm *shm;
static struct zwlr_layer_shell_v1 *layerShell;
static struct zwlr_screencopy_manager_v1 *scManager;
static struct wl_seat *seat;
static struct wl_pointer *pointer;
static struct wl_keyboard *keyboard;
static struct wl_output *output;

static struct xkb_context *xkb;
static struct xkb_keymap *keymap;
static struct xkb_state *xkbState;

static int pointerX, pointerY;
static bool dragging = false;
static int dragStartX = 0;
static int dragStartY = 0;

bool running = true;
static struct wl_surface *surface;
static struct zwlr_layer_surface_v1 *layerSurface;
static struct wl_egl_window *eglWindow;
static EGLSurface eglSurface;
static EGLContext eglContext;

static GLuint vertexShader, fragmentShader, shaderProgram;
static GLuint vao;
static GLuint texture;

static struct {
  int width, height, stride;
  uint8_t *data;
  bool ready;
} screenshot;

static struct {
  float zoom;
  float x, y;
} camera;
float flashlightRadius = 50.0f;
bool flashlightEnabled = false;
bool gridEnabled = false;

void noop() {}

void registry_add(void *data, struct wl_registry *registry, uint32_t name,
                  const char *interface, uint32_t version) {
  if (!strcmp(interface, "wl_compositor")) {
    compositor = wl_registry_bind(registry, name, &wl_compositor_interface, 1);
  } else if (!strcmp(interface, "zwlr_layer_shell_v1")) {
    layerShell =
        wl_registry_bind(registry, name, &zwlr_layer_shell_v1_interface, 1);
  } else if (!strcmp(interface, "wl_seat")) {
    seat = wl_registry_bind(registry, name, &wl_seat_interface, 1);
  } else if (!strcmp(interface, "zwlr_screencopy_manager_v1")) {
    scManager = wl_registry_bind(registry, name,
                                 &zwlr_screencopy_manager_v1_interface, 1);
  } else if (!strcmp(interface, "wl_output")) {
    output = wl_registry_bind(registry, name, &wl_output_interface, 1);
  } else if (!strcmp(interface, "wl_shm")) {
    shm = wl_registry_bind(registry, name, &wl_shm_interface, 1);
  }
}

static struct wl_registry_listener registry_listener = {.global = &registry_add,
                                                        .global_remove = &noop};

void layer_surface_configure(void *data, struct zwlr_layer_surface_v1 *surface,
                             uint32_t serial, uint32_t width, uint32_t height) {
  printf("Layer surface configure: %d %d\n", width, height);
  zwlr_layer_surface_v1_ack_configure(surface, serial);
  wl_egl_window_resize(eglWindow, width, height, 0, 0);
};

static struct zwlr_layer_surface_v1_listener layer_surface_listener = {
    .configure = &layer_surface_configure, .closed = &noop};

void pointer_motion(void *data, struct wl_pointer *pointer, uint32_t time,
                    wl_fixed_t surface_x, wl_fixed_t surface_y) {
  pointerX = wl_fixed_to_int(surface_x);
  pointerY = wl_fixed_to_int(surface_y);

  if (dragging) {
    float deltaX = (dragStartX / camera.zoom) - (pointerX / camera.zoom);
    float deltaY = (dragStartY / camera.zoom) - (pointerY / camera.zoom);
    dragStartX = pointerX;
    dragStartY = pointerY;

    camera.x += deltaX;
    camera.y += deltaY;
  }
}

void pointer_axis(void *data, struct wl_pointer *pointer, uint32_t time,
                  uint32_t axis, wl_fixed_t value) {
  float scroll = wl_fixed_to_double(value);

  if (xkb_state_mod_name_is_active(xkbState, XKB_MOD_NAME_CTRL,
                                   XKB_STATE_MODS_EFFECTIVE)) {

    flashlightRadius -= scroll;
    if (flashlightRadius < 0)
      flashlightRadius = 0;
  } else {
    camera.zoom -= scroll / 150.0f * 8.0f;
    if (camera.zoom < 1)
      camera.zoom = 1;
  }
}

void pointer_button(void *data, struct wl_pointer *pointer, uint32_t serial,
                    uint32_t time, uint32_t button, uint32_t state) {
  if (button == BTN_LEFT) {
    if (!dragging) {
      dragStartX = pointerX;
      dragStartY = pointerY;
    }
    dragging = state == WL_POINTER_BUTTON_STATE_PRESSED;
  }
}

static struct wl_pointer_listener pointer_listener = {.enter = &noop,
                                                      .leave = &noop,
                                                      .motion = &pointer_motion,
                                                      .button = &pointer_button,
                                                      .axis = &pointer_axis};

void screencopy_buffer(void *data, struct zwlr_screencopy_frame_v1 *frame,
                       uint32_t format, uint32_t width, uint32_t height,
                       uint32_t stride) {
  size_t size = stride * height;

  int fd = memfd_create("screencopy-shm", 0);
  ftruncate(fd, size);

  screenshot.data = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
  screenshot.width = width;
  screenshot.height = height;

  struct wl_shm_pool *pool = wl_shm_create_pool(shm, fd, size);
  struct wl_buffer *buffer = wl_shm_pool_create_buffer(
      pool, 0, width, height, stride, WL_SHM_FORMAT_XRGB8888);

  zwlr_screencopy_frame_v1_copy(frame, buffer);
}

void screencopy_ready(void *data, struct zwlr_screencopy_frame_v1 *frame,
                      uint32_t tv_sec_hi, uint32_t tv_sec_lo,
                      uint32_t tv_nsec) {
  printf("Frame done!\n");
  screenshot.ready = true;
}

static struct zwlr_screencopy_frame_v1_listener screencopy_listener = {
    .buffer = &screencopy_buffer,
    .flags = &noop,
    .ready = &screencopy_ready,
    .failed = &noop // TODO: error if failed
};

void keyboard_keymap(void *data, struct wl_keyboard *keyboard, uint32_t format,
                     int32_t fd, uint32_t size) {
  void *buf;

  buf = mmap(NULL, size, PROT_READ, MAP_SHARED, fd, 0);

  keymap = xkb_keymap_new_from_string(xkb, buf, XKB_KEYMAP_FORMAT_TEXT_V1,
                                      XKB_KEYMAP_COMPILE_NO_FLAGS);
  assert(keymap != NULL);

  xkbState = xkb_state_new(keymap);
}

void keyboard_key(void *data, struct wl_keyboard *keyboad, uint32_t serial,
                  uint32_t time, uint32_t key, uint32_t state) {
  xkb_keycode_t keycode = key + 8;
  printf("key: %d\n", keycode);

  xkb_keysym_t keysym = xkb_state_key_get_one_sym(xkbState, keycode);

  if (state == WL_KEYBOARD_KEY_STATE_PRESSED) {
    if (keysym == XKB_KEY_Escape || keysym == XKB_KEY_q) {
      running = false;
    } else if (keysym == XKB_KEY_f) {
      flashlightEnabled = !flashlightEnabled;
    } else if (keysym == XKB_KEY_g) {
      gridEnabled = !gridEnabled;
    }
  }
}

void keyboard_modifiers(void *data, struct wl_keyboard *keyboard,
                        uint32_t serial, uint32_t mods_depressed,
                        uint32_t mods_latched, uint32_t mods_locked,
                        uint32_t group) {
  xkb_state_update_mask(xkbState, mods_depressed, mods_latched, mods_locked, 0,
                        0, group);
}

static struct wl_keyboard_listener keyboard_listener = {
    .keymap = &keyboard_keymap,
    .enter = &noop,
    .leave = &noop,
    .key = &keyboard_key,
    .modifiers = &keyboard_modifiers,
    .repeat_info = &noop};

void gl_debug_callback(GLenum source, GLenum type, GLuint id, GLenum severity,
                       GLsizei length, const GLchar *message,
                       const void *userParam) {
  fprintf(stderr, "GL: %s\n", message);
}

void create_window() {
  eglDisplay = eglGetDisplay(display);
  eglInitialize(eglDisplay, NULL, NULL);
  eglBindAPI(EGL_OPENGL_API);
  EGLint attributes[] = {EGL_RED_SIZE,  8, EGL_GREEN_SIZE, 8,
                         EGL_BLUE_SIZE, 8, EGL_NONE};
  EGLConfig config;
  EGLint num_config;
  eglChooseConfig(eglDisplay, attributes, &config, 1, &num_config);
  eglContext = eglCreateContext(eglDisplay, config, EGL_NO_CONTEXT, NULL);

  surface = wl_compositor_create_surface(compositor);

  eglWindow = wl_egl_window_create(surface, 1920, 1080);
  eglSurface = eglCreateWindowSurface(eglDisplay, config,
                                      (EGLNativeWindowType)eglWindow, NULL);
  eglMakeCurrent(eglDisplay, eglSurface, eglSurface, eglContext);

  layerSurface = zwlr_layer_shell_v1_get_layer_surface(
      layerShell, surface, NULL, ZWLR_LAYER_SHELL_V1_LAYER_OVERLAY, "doomer");
  zwlr_layer_surface_v1_set_anchor(layerSurface,
                                   ZWLR_LAYER_SURFACE_V1_ANCHOR_TOP |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_BOTTOM |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_LEFT |
                                       ZWLR_LAYER_SURFACE_V1_ANCHOR_RIGHT);
  zwlr_layer_surface_v1_set_exclusive_zone(layerSurface, -1);
  zwlr_layer_surface_v1_set_keyboard_interactivity(
      layerSurface, ZWLR_LAYER_SURFACE_V1_KEYBOARD_INTERACTIVITY_EXCLUSIVE);
  zwlr_layer_surface_v1_add_listener(layerSurface, &layer_surface_listener,
                                     NULL);
  wl_surface_commit(surface);
  wl_display_roundtrip(display);

  GLenum err = glewInit();
  if (err != GLEW_OK) {
    fprintf(stderr, "Error: %s\n", glewGetErrorString(err));
  }
}

void compile_shader(unsigned int shader, const char *code, int size) {
  glShaderSource(shader, 1, &code, &size);
  glCompileShader(shader);

  int success;
  glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
  if (!success) {
    char infoLog[512];
    glGetShaderInfoLog(shader, 512, NULL, infoLog);
    fprintf(stderr, "Shader compilation error: %s\n", infoLog);
    exit(1);
  }
}

void init_shaders() {
  vertexShader = glCreateShader(GL_VERTEX_SHADER);
  fragmentShader = glCreateShader(GL_FRAGMENT_SHADER);

  compile_shader(vertexShader, vertexShaderSource, strlen(vertexShaderSource));
  compile_shader(fragmentShader, fragmentShaderSource,
                 strlen(fragmentShaderSource));

  shaderProgram = glCreateProgram();
  glAttachShader(shaderProgram, vertexShader);
  glAttachShader(shaderProgram, fragmentShader);
  glLinkProgram(shaderProgram);

  int success;
  glGetProgramiv(shaderProgram, GL_LINK_STATUS, &success);
  if (!success) {
    char infoLog[512];
    glGetProgramInfoLog(shaderProgram, 512, NULL, infoLog);
    fprintf(stderr, "Shader linking error: %s\n", infoLog);
    exit(1);
  }

  glUseProgram(shaderProgram);
}

void draw_window() {
  glClearColor(0, 0, 0, 1);
  glClear(GL_COLOR_BUFFER_BIT);

  glUseProgram(shaderProgram);
  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D, texture);
  glUniform1f(glGetUniformLocation(shaderProgram, "cameraScale"), camera.zoom);
  glUniform2f(glGetUniformLocation(shaderProgram, "cameraPos"), camera.x,
              camera.y);
  glUniform1f(glGetUniformLocation(shaderProgram, "flRadius"),
              flashlightRadius);
  glUniform1f(glGetUniformLocation(shaderProgram, "flShadow"),
              flashlightEnabled ? 0.7 : 0.0);
  glUniform1i(glGetUniformLocation(shaderProgram, "gridEnabled"), gridEnabled);
  glUniform2f(glGetUniformLocation(shaderProgram, "windowSize"), 1920, 1080);
  glUniform2f(glGetUniformLocation(shaderProgram, "screenshotSize"),
              screenshot.width, screenshot.height);
  glUniform2f(glGetUniformLocation(shaderProgram, "cursorPos"), pointerX,
              pointerY);
  glBindVertexArray(vao);
  glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
  glDrawElements(GL_TRIANGLES, 6, GL_UNSIGNED_INT, NULL);
  glBindVertexArray(0);

  eglSwapBuffers(eglDisplay, eglSurface);
}

int main() {
  camera.zoom = 1.0f;

  xkb = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
  assert(xkb != NULL);

  display = wl_display_connect(NULL);
  assert(display != NULL);

  struct wl_registry *registry = wl_display_get_registry(display);
  wl_registry_add_listener(registry, &registry_listener, NULL);
  wl_display_roundtrip(display);

  pointer = wl_seat_get_pointer(seat);
  keyboard = wl_seat_get_keyboard(seat);
  wl_keyboard_add_listener(keyboard, &keyboard_listener, NULL);
  wl_pointer_add_listener(pointer, &pointer_listener, NULL);

  // Start output capture
  struct zwlr_screencopy_frame_v1 *screencopyFrame =
      zwlr_screencopy_manager_v1_capture_output(scManager, 0, output);
  zwlr_screencopy_frame_v1_add_listener(screencopyFrame, &screencopy_listener,
                                        NULL);

  create_window();
  init_shaders();

  // Enable OpenGL debug printing
  glDebugMessageCallback(gl_debug_callback, NULL);
  glEnable(GL_DEBUG_OUTPUT);

  // Wait for the screenshot to be ready
  while (!screenshot.ready) {
    printf("wait...\n");
    wl_display_dispatch(display);
  }

  float w = 1920, h = 1080;
  // x y z u v
  GLfloat vertices[] = {
      w, 0, 0.0, 1.0, 1.0, // Top right
      w, h, 0.0, 1.0, 0.0, // Bottom right
      0, h, 0.0, 0.0, 0.0, // Bottom left
      0, 0, 0.0, 0.0, 1.0  // Top left
  };
  GLuint indices[] = {0, 1, 3, 1, 2, 3};

  GLuint vbo, ebo;
  glGenVertexArrays(1, &vao);
  glGenBuffers(1, &vbo);
  glGenBuffers(1, &ebo);

  glBindVertexArray(vao);
  // VBO for vertices
  glBindBuffer(GL_ARRAY_BUFFER, vbo);
  glBufferData(GL_ARRAY_BUFFER, sizeof(vertices), vertices, GL_STATIC_DRAW);

  // EBO for indices
  glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ebo);
  glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(indices), indices,
               GL_STATIC_DRAW);

  // x y z u v
  GLsizei stride = 5 * sizeof(GLfloat);

  // Position
  glVertexAttribPointer(0, 3, GL_FLOAT, false, stride, NULL);
  glEnableVertexAttribArray(0);

  // UV
  glVertexAttribPointer(1, 2, GL_FLOAT, false, stride,
                        (const void *)(3 * sizeof(GLfloat)));
  glEnableVertexAttribArray(1);

  glEnable(GL_TEXTURE_2D);

  // Convert texture from XRGB to RGB
  uint8_t converted[screenshot.width * screenshot.height * 3];
  for (size_t i = 0; i < (screenshot.width * screenshot.height); i++) {
    converted[3 * i + 0] = screenshot.data[4 * i + 2];
    converted[3 * i + 1] = screenshot.data[4 * i + 1];
    converted[3 * i + 2] = screenshot.data[4 * i + 0];
  }

  // Create the texture
  glGenTextures(1, &texture);
  glActiveTexture(GL_TEXTURE0);
  glBindTexture(GL_TEXTURE_2D, texture);
  glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, screenshot.width, screenshot.height, 0,
               GL_RGB, GL_UNSIGNED_BYTE, converted);
  glGenerateMipmap(GL_TEXTURE_2D);

  glBindTexture(GL_TEXTURE_2D, texture);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_BORDER);
  glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_BORDER);

  while (running) {
    wl_display_dispatch_pending(display);
    draw_window();
  }
}