#include <X11/XKBlib.h>
#include <X11/Xlib.h>
#include <X11/keysym.h>
#include <errno.h>
#include <pwd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <unistd.h>

#define MAX_BINDS 256
#define MAX_CMD_LEN 512
#define SOCK_PATH "/tmp/piss.sock"

typedef struct {
  unsigned int mod;
  KeyCode keycode;
  char cmd[MAX_CMD_LEN];
} Binding;

static Display *dpy;
static Window root;
static Binding bindings[MAX_BINDS];
static int bind_count = 0;

unsigned int parse_modifiers(const char *str) {
  unsigned int mod = 0;
  if (strcasestr(str, "super") || strcasestr(str, "mod4"))
    mod |= Mod4Mask;
  if (strcasestr(str, "shift"))
    mod |= ShiftMask;
  if (strcasestr(str, "ctrl") || strcasestr(str, "control"))
    mod |= ControlMask;
  if (strcasestr(str, "alt") || strcasestr(str, "mod1"))
    mod |= Mod1Mask;
  return mod;
}

void grab_key(unsigned int mod, KeyCode kc) {
  unsigned int numlock = Mod2Mask, capslock = LockMask;
  XGrabKey(dpy, kc, mod, root, True, GrabModeAsync, GrabModeAsync);
  XGrabKey(dpy, kc, mod | numlock, root, True, GrabModeAsync, GrabModeAsync);
  XGrabKey(dpy, kc, mod | capslock, root, True, GrabModeAsync, GrabModeAsync);
  XGrabKey(dpy, kc, mod | numlock | capslock, root, True, GrabModeAsync,
           GrabModeAsync);
}

void add_binding(char *msg) {
  if (bind_count >= MAX_BINDS)
    return;
  char *sep = strstr(msg, "|||");
  if (!sep)
    return;
  *sep = '\0';
  char *combo = msg, *cmd = sep + 3;

  char *last_plus = strrchr(combo, '+');
  if (!last_plus)
    return;
  *last_plus = '\0';
  char *key_name = last_plus + 1;
  while (*key_name == ' ')
    key_name++;

  KeySym sym = XStringToKeysym(key_name);
  if (sym == NoSymbol)
    return;

  bindings[bind_count].mod = parse_modifiers(combo);
  bindings[bind_count].keycode = XKeysymToKeycode(dpy, sym);
  strncpy(bindings[bind_count].cmd, cmd, MAX_CMD_LEN - 1);

  grab_key(bindings[bind_count].mod, bindings[bind_count].keycode);
  bind_count++;
  XSync(dpy, False);
}

int client_mode(int retry, char *combo, char *cmd) {
  struct sockaddr_un addr = {.sun_family = AF_UNIX};
  strncpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path) - 1);

  while (1) {
    int sock = socket(AF_UNIX, SOCK_STREAM, 0);
    if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) == 0) {
      if (combo && cmd) {
        char buf[MAX_CMD_LEN * 2];
        snprintf(buf, sizeof(buf), "%s|||%s", combo, cmd);
        write(sock, buf, strlen(buf));
      }
      close(sock);
      return 0;
    }
    close(sock);
    if (!retry)
      return 1;
    usleep(50000);
  }
}

int make_executable_if_needed(char *cfg) {
  struct stat st;
  if (stat(cfg, &st) == 0) {
    if (!(st.st_mode & S_IXUSR)) {
      if (chmod(cfg, st.st_mode | S_IXUSR) == -1) {
        perror("chmod");
        return -1;
      }
    }
  } else {
    perror("stat");
    return -1;
  }
  return 0;
}

int main(int argc, char *argv[]) {
  if (argc >= 2 && strcmp(argv[1], "wait") == 0)
    return client_mode(1, NULL, NULL);
  if (argc >= 4 && strcmp(argv[1], "bind") == 0)
    return client_mode(0, argv[2], argv[3]);

  dpy = XOpenDisplay(NULL);
  if (!dpy)
    return 1;
  root = DefaultRootWindow(dpy);
  XSelectInput(dpy, root, KeyPressMask);

  int serv = socket(AF_UNIX, SOCK_STREAM, 0);
  struct sockaddr_un addr = {.sun_family = AF_UNIX};
  strncpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path) - 1);
  unlink(SOCK_PATH);
  bind(serv, (struct sockaddr *)&addr, sizeof(addr));
  listen(serv, 10);

  char path[512];
  ssize_t len = readlink("/proc/self/exe", path, sizeof(path) - 1);
  if (len != -1) {
    path[len] = '\0';
    setenv("PISS_BIN", path, 1);
  }

  if (fork() == 0) {
    char *cfg = (argc > 1) ? argv[1] : NULL;
    if (!cfg) {
      static char def[512];
      snprintf(def, 511, "%s/.config/piss/pissrc", getpwuid(getuid())->pw_dir);
      cfg = def;
    }

    if (make_executable_if_needed(cfg) == -1) {
      exit(1);
    }

    char *args[] = {cfg, NULL};
    execvp(cfg, args);
    perror("execvp");
    exit(1);
  }

  int xfd = ConnectionNumber(dpy);
  while (1) {
    fd_set fds;
    FD_ZERO(&fds);
    FD_SET(xfd, &fds);
    FD_SET(serv, &fds);
    int max = (xfd > serv) ? xfd : serv;

    if (select(max + 1, &fds, NULL, NULL, NULL) > 0) {
      if (FD_ISSET(serv, &fds)) {
        int cli = accept(serv, NULL, NULL);
        char buf[MAX_CMD_LEN * 2] = {0};
        if (read(cli, buf, sizeof(buf) - 1) > 0)
          add_binding(buf);
        close(cli);
      }
      while (XPending(dpy)) {
        XEvent ev;
        XNextEvent(dpy, &ev);
        if (ev.type == KeyPress) {
          unsigned int state = ev.xkey.state & ~(Mod2Mask | LockMask);
          for (int i = 0; i < bind_count; i++) {
            if (bindings[i].keycode == ev.xkey.keycode &&
                bindings[i].mod == state) {
              if (fork() == 0) {
                setsid();
                execl("/bin/sh", "sh", "-c", bindings[i].cmd, NULL);
                exit(0);
              }
            }
          }
        }
      }
    }
  }
  return 0;
}