#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L

#include "waymini.h"

#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>

#include <linux/memfd.h>
#include <wayland-client.h>

#include "xdg-shell-client-protocol.h"

struct shm_state {
  int fd;
  size_t size;
  void *map;
  struct wl_shm *shm_iface;
  struct wl_buffer *wl_buffer;
};

struct window_state {
  struct wl_compositor *compositor;
  struct wl_surface *surface;
  struct xdg_wm_base *wm_base;
  struct xdg_surface *xdg_surface;
  struct xdg_toplevel *xdg_toplevel;

  uint32_t width, height;
  uint32_t bpp;
  uint32_t stride;
  int configured;
};

struct input_state {
  struct wl_seat *seat;
  struct wl_keyboard *keyboard;

  waymini_key_cb on_keycode;
  void *user;
};

struct waymini {
  struct wl_display *display;
  struct wl_registry *registry;

  struct wl_compositor *compositor;
  struct wl_shm *shm_iface;
  struct xdg_wm_base *wm_base;

  struct shm_state shm;
  struct window_state win;
  struct input_state in;

  int should_close;
};

static void on_xdg_surface_configure(void *data, struct xdg_surface *xdg_surface, uint32_t serial) {
  struct waymini *wm = data;
  xdg_surface_ack_configure(xdg_surface, serial);
  wm->win.configured = 1;
}

static const struct xdg_surface_listener xdg_surface_listener = {
  .configure = on_xdg_surface_configure,
};

static void on_toplevel_configure(void *data, struct xdg_toplevel *toplevel,
                                  int32_t width, int32_t height,
                                  struct wl_array *states) {
  (void)data; (void)toplevel; (void)width; (void)height; (void)states;
}

static void on_toplevel_close(void *data, struct xdg_toplevel *toplevel) {
  (void)toplevel;
  struct waymini *wm = data;
  wm->should_close = 1;
}

static const struct xdg_toplevel_listener xdg_toplevel_listener = {
  .configure = on_toplevel_configure,
  .close = on_toplevel_close,
};

static void on_wm_base_ping(void *data, struct xdg_wm_base *wm_base, uint32_t serial) {
  (void)data;
  xdg_wm_base_pong(wm_base, serial);
}

static const struct xdg_wm_base_listener wm_base_listener = {
  .ping = on_wm_base_ping,
};

static void on_keymap(void *data, struct wl_keyboard *keyboard, uint32_t format, int fd, uint32_t size) {
  (void)keyboard;
  struct input_state *in = data;

  if (format != WL_KEYBOARD_KEYMAP_FORMAT_XKB_V1) {
    close(fd);
    return;
  }
  char *buf = malloc(size);
  if (buf) {
    size_t off = 0;
    while (off < size) {
      ssize_t r = read(fd, buf + off, size - off);
      if (r <= 0) break;
      off += (size_t)r;
    }
    free(buf);
  } else {
    (void)read(fd, buf, 0);
  }
  close(fd);
  (void)in;
}

static void on_keyboard_key(void *data, struct wl_keyboard *keyboard, uint32_t serial,
                             uint32_t time, uint32_t key, uint32_t state) {
  (void)keyboard; (void)serial; (void)time;
  struct input_state *in = data;
  if (in->on_keycode) in->on_keycode(in->user, key, state);
}

static void on_keyboard_enter(void *data, struct wl_keyboard *keyboard,
                              uint32_t serial, struct wl_surface *surface,
                              struct wl_array *keys) {
  (void)data; (void)keyboard; (void)serial; (void)surface; (void)keys;
}

static void on_keyboard_leave(void *data, struct wl_keyboard *keyboard,
                              uint32_t serial, struct wl_surface *surface) {
  (void)data; (void)keyboard; (void)serial; (void)surface;
}

static void on_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) {
  (void)data; (void)keyboard; (void)serial;
  (void)mods_depressed; (void)mods_latched; (void)mods_locked; (void)group;
}

static void on_keyboard_repeat_info(void *data, struct wl_keyboard *keyboard,
                                    int32_t rate, int32_t delay) {
  (void)data; (void)keyboard; (void)rate; (void)delay;
}

static const struct wl_keyboard_listener keyboard_listener = {
  .keymap = on_keymap,
  .enter = on_keyboard_enter,
  .leave = on_keyboard_leave,
  .key = on_keyboard_key,
  .modifiers = on_keyboard_modifiers,
  .repeat_info = on_keyboard_repeat_info,
};

static void on_seat_capabilities(void *data, struct wl_seat *seat, enum wl_seat_capability caps) {
  struct input_state *in = data;
  if ((caps & WL_SEAT_CAPABILITY_KEYBOARD) && !in->keyboard) {
    in->seat = seat;
    in->keyboard = wl_seat_get_keyboard(seat);
    wl_keyboard_add_listener(in->keyboard, &keyboard_listener, in);
  }
}

static void on_seat_name(void *data, struct wl_seat *seat, const char *name) {
  (void)data; (void)seat; (void)name;
}

static const struct wl_seat_listener seat_listener = {
  .capabilities = on_seat_capabilities,
  .name = on_seat_name,
};

static void on_registry_global(void *data, struct wl_registry *registry,
                               uint32_t name, const char *interface, uint32_t version) {
  struct waymini *wm = data;
  (void)registry;

  if (strcmp(interface, "wl_compositor") == 0) {
    wm->compositor = wl_registry_bind(wm->registry, name, &wl_compositor_interface, 4);
  } else if (strcmp(interface, "wl_shm") == 0) {
    wm->shm_iface = wl_registry_bind(wm->registry, name, &wl_shm_interface, 1);
    wm->shm.shm_iface = wm->shm_iface;
  } else if (strcmp(interface, "xdg_wm_base") == 0) {
    wm->wm_base = wl_registry_bind(wm->registry, name, &xdg_wm_base_interface, 1);
    xdg_wm_base_add_listener(wm->wm_base, &wm_base_listener, wm);
  } else if (strcmp(interface, "wl_seat") == 0) {
    struct wl_seat *seat = wl_registry_bind(wm->registry, name, &wl_seat_interface, version);
    wl_seat_add_listener(seat, &seat_listener, &wm->in);
  }
}

static void on_registry_global_remove(void *data, struct wl_registry *registry, uint32_t name) {
  (void)data; (void)registry; (void)name;
}

static const struct wl_registry_listener registry_listener = {
  .global = on_registry_global,
  .global_remove = on_registry_global_remove,
};

static void waymini_destroy_internal(struct waymini *wm);

static int create_shm_buffer(struct waymini *wm) {
  struct shm_state *s = &wm->shm;
  struct window_state *w = &wm->win;

  s->size = (size_t)w->stride * (size_t)w->height;
  s->fd = memfd_create("waymini-shm", MFD_CLOEXEC);
  if (s->fd < 0) { perror("memfd_create"); return -1; }

  if (ftruncate(s->fd, (off_t)s->size) < 0) { perror("ftruncate"); return -1; }

  s->map = mmap(NULL, s->size, PROT_READ | PROT_WRITE, MAP_SHARED, s->fd, 0);
  if (s->map == MAP_FAILED) { perror("mmap"); return -1; }

  struct wl_buffer *buf;
  struct wl_shm_pool *pool = wl_shm_create_pool(s->shm_iface, s->fd, (int)s->size);
  buf = wl_shm_pool_create_buffer(pool, 0, w->width, w->height, (int)w->stride, WL_SHM_FORMAT_XRGB8888);
  wl_shm_pool_destroy(pool);
  if (!buf) { fprintf(stderr, "wl_shm_pool_create_buffer failed\n"); return -1; }

  s->wl_buffer = buf;
  return 0;
}

struct waymini *waymini_create(uint32_t width, uint32_t height,
                               waymini_key_cb key_cb, void *user) {
  struct waymini *wm = calloc(1, sizeof(*wm));
  if (!wm) return NULL;

  wm->display = wl_display_connect(NULL);
  if (!wm->display) { fprintf(stderr, "wl_display_connect failed\n"); free(wm); return NULL; }

  wm->registry = wl_display_get_registry(wm->display);
  wl_registry_add_listener(wm->registry, &registry_listener, wm);
  wl_display_roundtrip(wm->display);

  if (!wm->compositor || !wm->shm_iface || !wm->wm_base) {
    fprintf(stderr, "missing globals\n");
    waymini_destroy_internal(wm);
    return NULL;
  }

  wm->win.compositor = wm->compositor;
  wm->win.surface = wl_compositor_create_surface(wm->compositor);

  wm->win.xdg_surface = xdg_wm_base_get_xdg_surface(wm->wm_base, wm->win.surface);
  xdg_surface_add_listener(wm->win.xdg_surface, &xdg_surface_listener, wm);

  wm->win.xdg_toplevel = xdg_surface_get_toplevel(wm->win.xdg_surface);
  xdg_toplevel_add_listener(wm->win.xdg_toplevel, &xdg_toplevel_listener, wm);
  xdg_toplevel_set_title(wm->win.xdg_toplevel, "waymini");
  xdg_toplevel_set_app_id(wm->win.xdg_toplevel, "waymini");

  wm->in.on_keycode = key_cb;
  wm->in.user = user;

  wm->win.width = width;
  wm->win.height = height;
  wm->win.bpp = 4;
  wm->win.stride = width * wm->win.bpp;

  wl_surface_commit(wm->win.surface);

  while (!wm->win.configured) {
    if (wl_display_dispatch(wm->display) < 0) {
      fprintf(stderr, "dispatch failed waiting for configure\n");
      waymini_destroy_internal(wm);
      return NULL;
    }
  }

  if (create_shm_buffer(wm) < 0) {
    fprintf(stderr, "create_shm_buffer failed\n");
    waymini_destroy_internal(wm);
    return NULL;
  }

  return wm;
}

uint32_t waymini_get_width(const struct waymini *wm) { return wm ? wm->win.width : 0; }
uint32_t waymini_get_height(const struct waymini *wm) { return wm ? wm->win.height : 0; }
uint32_t waymini_get_stride(const struct waymini *wm) { return wm ? wm->win.stride : 0; }
void    *waymini_get_pixels(const struct waymini *wm) { return wm ? wm->shm.map : NULL; }

void waymini_present(struct waymini *wm) {
  if (!wm || !wm->win.configured) return;
  wl_surface_attach(wm->win.surface, wm->shm.wl_buffer, 0, 0);
  wl_surface_commit(wm->win.surface);
  wl_display_flush(wm->display);
}

int waymini_dispatch(struct waymini *wm) {
  if (!wm || !wm->display) return -1;
  return wl_display_dispatch(wm->display);
}

int waymini_poll(struct waymini *wm) {
  if (!wm || !wm->display) return -1;
  wl_display_dispatch_pending(wm->display);
  wl_display_flush(wm->display);
  int ret = wl_display_prepare_read(wm->display);
  if (ret < 0) return wl_display_dispatch_pending(wm->display) >= 0 ? 1 : -1;

  int fd = wl_display_get_fd(wm->display);
  fd_set rfds;
  FD_ZERO(&rfds);
  FD_SET(fd, &rfds);
  struct timeval tv = { 0, 0 };
  ret = select(fd + 1, &rfds, NULL, NULL, &tv);
  if (ret > 0) {
    wl_display_read_events(wm->display);
    wl_display_dispatch_pending(wm->display);
    return 1;
  }
  wl_display_cancel_read(wm->display);
  return 0;
}

void waymini_request_close(struct waymini *wm) {
  if (wm) wm->should_close = 1;
}

int waymini_should_close(const struct waymini *wm) {
  return wm ? wm->should_close : 0;
}

static void waymini_destroy_internal(struct waymini *wm) {
  if (!wm) return;

  if (wm->in.keyboard) wl_keyboard_destroy(wm->in.keyboard);
  if (wm->in.seat) wl_seat_destroy(wm->in.seat);

  if (wm->shm.wl_buffer) wl_buffer_destroy(wm->shm.wl_buffer);
  if (wm->shm.map && wm->shm.size) munmap(wm->shm.map, wm->shm.size);
  if (wm->shm.fd >= 0) close(wm->shm.fd);

  if (wm->win.xdg_toplevel) xdg_toplevel_destroy(wm->win.xdg_toplevel);
  if (wm->win.xdg_surface) xdg_surface_destroy(wm->win.xdg_surface);
  if (wm->win.surface) wl_surface_destroy(wm->win.surface);

  if (wm->wm_base) xdg_wm_base_destroy(wm->wm_base);
  if (wm->compositor) wl_compositor_destroy(wm->compositor);
  if (wm->shm_iface) wl_shm_destroy(wm->shm_iface);

  if (wm->registry) wl_registry_destroy(wm->registry);
  if (wm->display) wl_display_disconnect(wm->display);
  free(wm);
}

void waymini_destroy(struct waymini *wm) {
  waymini_destroy_internal(wm);
}