#include <fcntl.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/time.h>
#include <sys/timerfd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <wayland-client.h>
#include <wayland-server.h>

#include "zwp-idle-inhibit-manager-v1-protocol.h"
#include "ext-idle-notify-v1-protocol.h"
#include "xdg-shell-client-protocol.h"

#define LEN(X) (sizeof(X) / sizeof(X[0]))

static struct wl_compositor *compositor;
static struct wl_display *display;
static struct wl_registry *registry;
static struct wl_surface *surface;
static struct wl_seat *seat;

static struct xdg_surface *xdgsurface;
static struct xdg_wm_base *base;
static struct xdg_toplevel *toplevel;

static struct zwp_idle_inhibitor_v1 *inhibitor;
static struct zwp_idle_inhibit_manager_v1 *inhibitmanager;

static struct ext_idle_notification_v1 *notif;
static struct ext_idle_notifier_v1 *notifier;

static int subcommand = -1;
enum { INSOMNAIC, INHIBITCHECK, PROTOCOLS };

static void die(const char *fmt, ...);
static void pong(void *data, struct xdg_wm_base *base, uint32_t serial);
static void reghandler(void *data, struct wl_registry *registry, uint32_t id,
  const char *interface, uint32_t version);
static void idling(void *data, struct ext_idle_notification_v1 *notif);
static void interrupt();

static void
die(const char *fmt, ...)
{
  va_list ap;

  va_start(ap, fmt);
  vfprintf(stderr, fmt, ap);
  va_end(ap);

  if (fmt[0] && fmt[strlen(fmt)-1] == ':') {
    fputc(' ', stderr);
    perror(NULL);
  } else
    fputc('\n', stderr);
  exit(1);
}

static void
pong(void *data, struct xdg_wm_base *base, uint32_t serial)
{
  xdg_wm_base_pong(base, serial);
}

static const struct xdg_wm_base_listener xdglistener = {
  .ping = pong,
};

static void
reghandler(void *data, struct wl_registry *registry, uint32_t id,
    const char *interface, uint32_t version)
{
  if (strcmp(interface, wl_compositor_interface.name) == 0) {
    compositor = wl_registry_bind(registry, id, &wl_compositor_interface, 1);
  } else if (strcmp(interface, xdg_wm_base_interface.name) == 0) {
    base = wl_registry_bind(registry, id, &xdg_wm_base_interface, 1);
    xdg_wm_base_add_listener(base, &xdglistener, NULL);
  } else if (strcmp(interface, wl_seat_interface.name) == 0) {
    seat = wl_registry_bind(registry, id, &wl_seat_interface, 7);
  } else if (strcmp(interface, zwp_idle_inhibit_manager_v1_interface.name) == 0) {
    if (subcommand == INSOMNAIC) {
      inhibitmanager = wl_registry_bind(registry, id,
        &zwp_idle_inhibit_manager_v1_interface, 1);
    }
  } else if (strcmp(interface, ext_idle_notifier_v1_interface.name) == 0) {
    if (subcommand == INHIBITCHECK) {
      notifier = wl_registry_bind(registry, id, &ext_idle_notifier_v1_interface, 1);
    }
  }
  if (subcommand == PROTOCOLS)
    fprintf(stdout, "%s (%d)\n", interface, version);
}

static const struct wl_registry_listener reglistener = {
  .global = reghandler,
};

static void
idling(void *data, struct ext_idle_notification_v1 *notif)
{
  printf("No idle inhibitor found :)\n");
  exit(0);
}

static void
interrupt()
{
  printf("Idle inhibitor found :(\n");
  exit(1);
}

static const struct ext_idle_notification_v1_listener idlelistener = {
  .idled = idling,
};

int
main(int argc, char *argv[])
{
  int i, c;
  struct timeval start, end;
  pid_t pid;

  while ((c = getopt(argc, argv, "hvf")) != -1) {
    switch (c) {
      case 'v': die("%s-%s", argv[0], VERSION); break;
      case 'f':
        if (fork() != 0)
          exit(0);
        break;
      case 'h':
      default: die("usage: %s [-hvf]", argv[0]); break;
    }
  }

  for (i = 0; i < argc; i++) {
    if (strcmp("insomniac", argv[i]) == 0) {
      fprintf(stderr, "Adding an idle lock on the compositor\n");
      subcommand = INSOMNAIC;
    } else if (strcmp("inhibitcheck", argv[i]) == 0) {
      fprintf(stderr, "Checking for an idle inhibitor\n");
      subcommand = INHIBITCHECK;
    } else if (strcmp("protocols", argv[i]) == 0) {
      fprintf(stderr, "Checking for protocols\n");
      subcommand = PROTOCOLS;
    } else if (argv[i] == argv[argc - 1]) {
      fprintf(stderr, "No valid command specified\n");
      exit(-1);
    }
  }

  /* get the parent process pid to make killing it easier later */
  pid = getpid();
  
  /* setup wayland stuff */
  display = wl_display_connect(NULL);
  if (!display)
    die("%s: failed to connect to wayland display :(", argv[0]);

  registry = wl_display_get_registry(display);
  wl_registry_add_listener(registry, &reglistener, NULL);
  wl_display_roundtrip(display);
  if (subcommand == PROTOCOLS)
    exit(0);
  if (!compositor || !seat)
    die("%s: compositor doesn't meet necesary protocols", argv[0]);
  if (!notifier && subcommand == INHIBITCHECK)
    die("%s: compositor doesn't support the idle notify protocol", argv[0]);
  if (!inhibitmanager && subcommand == INSOMNAIC)
    die("%s: compositor doesn't support the inhibit manager protocol", argv[0]);

  /* register event */
  if (subcommand == INHIBITCHECK) {
    notif = ext_idle_notifier_v1_get_idle_notification(notifier, 1, seat);
    ext_idle_notification_v1_add_listener(notif, &idlelistener, NULL);

    /* event should be registered now get starting time so we can end command
     * if we should've seen the idle event by now */
    gettimeofday(&start, NULL);
  }
  wl_display_roundtrip(display);

  surface = wl_compositor_create_surface(compositor);
  xdgsurface = xdg_wm_base_get_xdg_surface(base, surface);
  toplevel = xdg_surface_get_toplevel(xdgsurface);

  wl_display_roundtrip(display);
  wl_surface_commit(surface);

  /* register stuff */
  switch (subcommand) {
    case INSOMNAIC:
      inhibitor = zwp_idle_inhibit_manager_v1_create_inhibitor(inhibitmanager, surface);
      break;
    case INHIBITCHECK:
      /* check in a seperate thread so wayland doesn't block or get blocked */
      if (fork() == 0)
        while (1) {
          gettimeofday(&end, NULL);
          /* check if the program has been running for 10 milliseconds, and then
           * terminate because at this point we should've seen an idle event */
          if (end.tv_usec - start.tv_usec > 10000) {
            kill(pid, SIGPOLL);
            exit(1);
          }
        }

      /* listen for the SIGPOLL signal */
      signal(SIGPOLL, interrupt);
      break;
  }
  wl_display_roundtrip(display);

  /* start listening for wayland events */
  while (wl_display_dispatch(display) != -1);

  /* cleanup time */
  if (subcommand == INSOMNAIC)
    zwp_idle_inhibitor_v1_destroy(inhibitor);
  zwp_idle_inhibit_manager_v1_destroy(inhibitmanager);
  ext_idle_notifier_v1_destroy(notifier);
  wl_registry_destroy(registry);
  wl_display_destroy(display);
  wl_compositor_destroy(compositor);
}