#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <sys/mman.h>

#include <wayland-client.h>

#include "ext-image-capture-source-v1.prot.h"
#include "ext-image-copy-capture-v1.prot.h"

#include "rif.h"

#define ASSERT(cond, msg)\
  if(!(cond)) {\
    fprintf(stderr, "Runtime assertion %s at %s:%d failed: %s\n", #cond, __FILE__, __LINE__, msg);\
    exit(-1);\
  }

// Not included in any standard header
// Thus, declaring manually.
// See memfd_create(2).
int memfd_create(const char* name, unsigned int flags);

struct image_data {
  uint32_t width;
  uint32_t height;
  uint32_t shm_format;
};

struct wl_output* stream_output = NULL; 
struct wl_shm* wl_shm = NULL;
struct wl_buffer* frame_buffer = NULL;
struct ext_output_image_capture_source_manager_v1* output_source_manager = NULL;
struct ext_image_copy_capture_manager_v1* capture_manager = NULL;

void* shm_data = NULL;

struct image_data img_data;

uint8_t is_frame_ready = 0;
uint8_t is_capture_ready = 0;

uint32_t buf_size;

void handle_global(
    void* data, 
    struct wl_registry* registry, 
    uint32_t name,
    const char* interface, 
    uint32_t version) {
  if(strcmp(interface, wl_output_interface.name) == 0) {
    stream_output = wl_registry_bind(registry, name, &wl_output_interface, version);
  }
  if(strcmp(interface, wl_shm_interface.name) == 0) {
    wl_shm = wl_registry_bind(registry, name, &wl_shm_interface, version);
  }
  if(strcmp(interface, ext_output_image_capture_source_manager_v1_interface.name) == 0) {
    output_source_manager = wl_registry_bind(registry, name, &ext_output_image_capture_source_manager_v1_interface, version);
  }
  if(strcmp(interface, ext_image_copy_capture_manager_v1_interface.name) == 0) {
    capture_manager = wl_registry_bind(registry, name, &ext_image_copy_capture_manager_v1_interface, version);
  }
}

void handle_global_remove(
    void* data,
    struct wl_registry* registry,
    uint32_t name) {
	// Who cares
}

const struct wl_registry_listener reg_callbacks = {
	.global = handle_global,
	.global_remove = handle_global_remove,
};

void buffer_size(void *data,
        struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1,
        uint32_t width,
        uint32_t height) {
  printf("x: %d, y: %d\n", width, height);
  img_data.width = width;
  img_data.height = height;
}

void shm_format(void *data,
       struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1,
       uint32_t format) {
  img_data.shm_format = format;
  ASSERT(format == WL_SHM_FORMAT_XBGR8888 || format == WL_SHM_FORMAT_ABGR8888, "Unsupported buffer format.");
}

void dmabuf_device(void *data,
          struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1,
          struct wl_array *device) {

}


void dmabuf_format(void *data,
          struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1,
          uint32_t format,
          struct wl_array *modifiers) {

}

void done(void *data,
       struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1) {
  const int shm_fd = memfd_create("wl_shm data", 0);
  ASSERT(shm_fd != -1, "Failed to create memfd.");

  // Format is hardcoded to chan width 4
  const uint32_t chan_width = 4;
  const uint32_t stride = img_data.width * chan_width;
  buf_size = img_data.height * stride;

  ASSERT(ftruncate(shm_fd, buf_size) != -1, "Failed to ftruncate the memfd.");

  shm_data = mmap(NULL, buf_size, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0);
  ASSERT(shm_data != MAP_FAILED, "Failed to mmap the memfd.");

  struct wl_shm_pool* pool = wl_shm_create_pool(wl_shm, shm_fd, buf_size);
  ASSERT(pool != NULL, "Failed to create the wl_shm_pool.");

  frame_buffer = wl_shm_pool_create_buffer(pool, 0, img_data.width, img_data.height, stride, img_data.shm_format);
  ASSERT(frame_buffer != NULL, "Failed to create the wl_buffer.");

  is_frame_ready = 1;
}


void stopped(void *data,
    struct ext_image_copy_capture_session_v1 *ext_image_copy_capture_session_v1) {

}

const struct ext_image_copy_capture_session_v1_listener capture_listener = {
  .buffer_size = buffer_size,
  .shm_format = shm_format,
  .dmabuf_device = dmabuf_device,
  .dmabuf_format = dmabuf_format,
  .done = done,
  .stopped = stopped
};


void transform(void *data,
      struct ext_image_copy_capture_frame_v1 *ext_image_copy_capture_frame_v1,
      uint32_t transform) {

}
void damage(void *data,
         struct ext_image_copy_capture_frame_v1 *ext_image_copy_capture_frame_v1,
         int32_t x,
         int32_t y,
         int32_t width,
         int32_t height) {

}
void presentation_time(void *data,
        struct ext_image_copy_capture_frame_v1 *ext_image_copy_capture_frame_v1,
        uint32_t tv_sec_hi,
        uint32_t tv_sec_lo,
        uint32_t tv_nsec) {

}
void ready(void *data,
        struct ext_image_copy_capture_frame_v1 *ext_image_copy_capture_frame_v1) {
  is_capture_ready = 1;
}
void failed(void *data,
         struct ext_image_copy_capture_frame_v1 *ext_image_copy_capture_frame_v1,
         uint32_t reason) {
  switch(reason) {
    case 0:
      printf("Capture failed. Reason: Unknown Runtime Error.");
      break;
    case 1:
      printf("Capture failed. Reason: Buffer Constraints Mismatch.");
      break;
    case 2:
      printf("Capture failed. Reason: Session Is No Longer Available.");
      break;
  }
}

const struct ext_image_copy_capture_frame_v1_listener frame_listener = {
  .transform = transform,
  .damage = damage,
  .presentation_time = presentation_time,
  .ready = ready,
  .failed = failed
};


int main() {
  struct wl_display* dpy = wl_display_connect(NULL);
  ASSERT(dpy != NULL, "Unable to connect to Wayland");

  struct wl_registry* registry = wl_display_get_registry(dpy);
	wl_registry_add_listener(registry, &reg_callbacks, NULL);
	wl_display_roundtrip(dpy);

  ASSERT(wl_shm != NULL, "core.wl_shm not supported.");

  ASSERT(stream_output != NULL, "No outputs found!");

  ASSERT(output_source_manager != NULL, "ext-image-capture-source-v1.ext_output_image_capture_source_manager_v1 not supported.");
  ASSERT(capture_manager != NULL, "ext-image-copy-capture-v1.ext_image_copy_capture_manager_v1 not supported.");

  struct ext_image_capture_source_v1* capture_source = 
    ext_output_image_capture_source_manager_v1_create_source(output_source_manager, stream_output);
  ASSERT(capture_source != NULL, "Failed creating a source for the default output!")
    
  struct ext_image_copy_capture_session_v1* capture_session =
    ext_image_copy_capture_manager_v1_create_session(capture_manager, capture_source, EXT_IMAGE_COPY_CAPTURE_MANAGER_V1_OPTIONS_PAINT_CURSORS);
  ASSERT(capture_session != NULL, "Failed creating a session for the capture!");

  ext_image_copy_capture_session_v1_add_listener(capture_session, &capture_listener, NULL);
  wl_display_roundtrip(dpy);

  while(is_frame_ready == 0 && wl_display_dispatch(dpy) != -1) {
    // Empty
  }

  struct ext_image_copy_capture_frame_v1* frame = ext_image_copy_capture_session_v1_create_frame(capture_session);
  ASSERT(frame != NULL, "Failed to create frame!");

  ext_image_copy_capture_frame_v1_add_listener(frame, &frame_listener, NULL);

  ext_image_copy_capture_frame_v1_attach_buffer(frame, frame_buffer);
  ext_image_copy_capture_frame_v1_damage_buffer(frame, 0, 0, img_data.width, img_data.height);
  ext_image_copy_capture_frame_v1_capture(frame);

  wl_display_roundtrip(dpy);

  while(is_capture_ready == 0 && wl_display_dispatch(dpy) != -1) {
    // Empty
  }

  printf("Saving screenshot.\n");

  // void write_rif_little(char* path, uint32_t width, uint32_t size, uint8_t format, void* data) {
  // Thanks to endianess shenanigans, ABGR/XBGR is RGBA. 
  write_rif_little("out.rif", img_data.width, buf_size, RIF_FORMAT_R8G8B8A8, shm_data);

  ext_image_copy_capture_frame_v1_destroy(frame);

  return 0;
}