#include <X11/Xlib.h>
#include <X11/keysym.h>
#include <xcb/xcb.h>
#include <xcb/xproto.h>

#include <errno.h>
#include <fcntl.h>
#include <pwd.h>
#include <signal.h>
#include <stdarg.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/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <unistd.h>

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

#define MOD_SHIFT XCB_MOD_MASK_SHIFT
#define MOD_CAPS XCB_MOD_MASK_LOCK
#define MOD_CTRL XCB_MOD_MASK_CONTROL
#define MOD_ALT XCB_MOD_MASK_1
#define MOD_NUMLOCK XCB_MOD_MASK_2
#define MOD_SUPER XCB_MOD_MASK_4

#define CLEAN_MASK (~(MOD_NUMLOCK | MOD_CAPS))

typedef struct {
  uint16_t mod;
  xcb_keycode_t keycode;
  char cmd[MAX_CMD_LEN];
} Binding;

static xcb_connection_t *conn = NULL;
static xcb_window_t root = 0;

static xcb_get_keyboard_mapping_reply_t *kbmap = NULL;
static xcb_keycode_t kbmap_min = 0;
static xcb_keycode_t kbmap_max = 0;

static Binding bindings[MAX_BINDS];
static int bind_count = 0;

static volatile sig_atomic_t g_quit = 0;

static void die(const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  fprintf(stderr, "piss: fatal: ");
  vfprintf(stderr, fmt, ap);
  va_end(ap);
  fputc('\n', stderr);
  exit(1);
}

static void warn(const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  fprintf(stderr, "piss: warning: ");
  vfprintf(stderr, fmt, ap);
  va_end(ap);
  fputc('\n', stderr);
}

static void handle_signal(int sig) {
  (void)sig;
  g_quit = 1;
}

static void refresh_keymap(void) {
  free(kbmap);
  xcb_get_keyboard_mapping_cookie_t ck = xcb_get_keyboard_mapping(
      conn, kbmap_min, (uint8_t)(kbmap_max - kbmap_min + 1));
  xcb_generic_error_t *e = NULL;
  kbmap = xcb_get_keyboard_mapping_reply(conn, ck, &e);
  if (e) {
    int code = e->error_code;
    free(e);
    die("xcb_get_keyboard_mapping error %d", code);
  }
  if (!kbmap)
    die("xcb_get_keyboard_mapping returned NULL");
}

static xcb_keycode_t keysym_to_keycode(xcb_keysym_t sym) {
  if (!kbmap)
    return 0;

  const xcb_keysym_t *table = xcb_get_keyboard_mapping_keysyms(kbmap);
  int kpk = kbmap->keysyms_per_keycode;

  for (int kc = kbmap_min; kc <= kbmap_max; kc++) {
    const xcb_keysym_t *row = table + (kc - kbmap_min) * kpk;
    for (int col = 0; col < kpk; col++) {
      if (row[col] == sym)
        return (xcb_keycode_t)kc;
    }
  }
  return 0;
}

static uint16_t parse_modifiers(const char *str) {
  uint16_t mod = 0;
  if (strcasestr(str, "super") || strcasestr(str, "mod4"))
    mod |= MOD_SUPER;
  if (strcasestr(str, "shift"))
    mod |= MOD_SHIFT;
  if (strcasestr(str, "ctrl") || strcasestr(str, "control"))
    mod |= MOD_CTRL;
  if (strcasestr(str, "alt") || strcasestr(str, "mod1"))
    mod |= MOD_ALT;
  return mod;
}

static int grab_key(uint16_t mod, xcb_keycode_t kc) {
  uint16_t mods[4] = {
      mod,
      (uint16_t)(mod | MOD_NUMLOCK),
      (uint16_t)(mod | MOD_CAPS),
      (uint16_t)(mod | MOD_NUMLOCK | MOD_CAPS),
  };
  int failed = 0;
  for (int i = 0; i < 4; i++) {
    xcb_void_cookie_t ck = xcb_grab_key_checked(
        conn, 1, root, mods[i], kc, XCB_GRAB_MODE_ASYNC, XCB_GRAB_MODE_ASYNC);
    xcb_generic_error_t *e = xcb_request_check(conn, ck);
    if (e) {
      if (e->error_code == XCB_ACCESS)
        warn("keycode %u mod 0x%04x already grabbed by another client"
             " — is another hotkey daemon running?",
             kc, mods[i]);
      else
        warn("XGrabKey error %d for keycode %u mod 0x%04x", e->error_code, kc,
             mods[i]);
      free(e);
      failed = 1;
    }
  }
  return failed ? -1 : 0;
}

static void ungrab_all(void) {
  for (int i = 0; i < bind_count; i++) {
    uint16_t m = bindings[i].mod;
    xcb_keycode_t kc = bindings[i].keycode;
    uint16_t mods[4] = {
        m,
        (uint16_t)(m | MOD_NUMLOCK),
        (uint16_t)(m | MOD_CAPS),
        (uint16_t)(m | MOD_NUMLOCK | MOD_CAPS),
    };
    for (int j = 0; j < 4; j++)
      xcb_ungrab_key(conn, kc, root, mods[j]);
  }
  xcb_flush(conn);
}

static int add_binding(char *msg) {
  if (bind_count >= MAX_BINDS) {
    warn("MAX_BINDS (%d) reached — ignoring", MAX_BINDS);
    return -1;
  }

  char *sep = strstr(msg, "|||");
  if (!sep) {
    warn("malformed bind message (no |||): %.80s", msg);
    return -1;
  }
  *sep = '\0';
  char *combo = msg;
  char *cmd = sep + 3;

  if (!*cmd) {
    warn("empty command for combo '%s'", combo);
    return -1;
  }

  char *last_plus = strrchr(combo, '+');
  if (!last_plus) {
    warn("no '+' in combo '%s'", combo);
    return -1;
  }
  *last_plus = '\0';
  char *keyname = last_plus + 1;

  while (*keyname == ' ' || *keyname == '\t')
    keyname++;
  {
    char *e = keyname + strlen(keyname);
    while (e > keyname && (e[-1] == ' ' || e[-1] == '\t'))
      *--e = '\0';
  }

  if (!*keyname) {
    warn("empty key name in combo");
    return -1;
  }

  KeySym sym = XStringToKeysym(keyname);
  if (sym == NoSymbol) {
    warn("unknown keysym '%s'", keyname);
    return -1;
  }

  xcb_keycode_t kc = keysym_to_keycode((xcb_keysym_t)sym);
  if (kc == 0) {
    warn("no keycode for keysym '%s' (0x%lx)", keyname, sym);
    return -1;
  }

  uint16_t mod = parse_modifiers(combo);

  for (int i = 0; i < bind_count; i++) {
    if (bindings[i].keycode == kc && bindings[i].mod == mod) {
      warn("overwriting existing binding for '%s+%s'", combo, keyname);
      strncpy(bindings[i].cmd, cmd, MAX_CMD_LEN - 1);
      bindings[i].cmd[MAX_CMD_LEN - 1] = '\0';
      return 0;
    }
  }

  bindings[bind_count].mod = mod;
  bindings[bind_count].keycode = kc;
  strncpy(bindings[bind_count].cmd, cmd, MAX_CMD_LEN - 1);
  bindings[bind_count].cmd[MAX_CMD_LEN - 1] = '\0';

  int grabbed = grab_key(mod, kc);
  if (grabbed == -1)
    warn("stored binding but grab failed for '%s+%s'", combo, keyname);

  bind_count++;
  xcb_flush(conn);
  fprintf(stderr, "piss: bound '%s+%s' → %s\n", combo, keyname, cmd);
  return grabbed;
}

static void spawn(const char *cmd) {
  pid_t pid = fork();
  if (pid < 0) {
    warn("fork: %s", strerror(errno));
    return;
  }
  if (pid == 0) {
    setsid();
    signal(SIGCHLD, SIG_DFL);
    long maxfd = sysconf(_SC_OPEN_MAX);
    if (maxfd < 0)
      maxfd = 1024;
    for (long fd = 3; fd < maxfd; fd++)
      close((int)fd);
    execl("/bin/sh", "sh", "-c", cmd, NULL);
    fprintf(stderr, "piss: execl: %s\n", strerror(errno));
    _exit(1);
  }
}

static int make_server_socket(void) {
  int fd = socket(AF_UNIX, SOCK_STREAM, 0);
  if (fd < 0)
    die("socket: %s", strerror(errno));

  if (fcntl(fd, F_SETFD, FD_CLOEXEC) < 0)
    die("fcntl(CLOEXEC) on server socket: %s", strerror(errno));

  struct sockaddr_un addr;
  memset(&addr, 0, sizeof(addr));
  addr.sun_family = AF_UNIX;
  strncpy(addr.sun_path, SOCK_PATH, sizeof(addr.sun_path) - 1);

  unlink(SOCK_PATH);

  if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0)
    die("bind(%s): %s", SOCK_PATH, strerror(errno));

  if (listen(fd, 32) < 0)
    die("listen: %s", strerror(errno));

  if (chmod(SOCK_PATH, 0600) < 0)
    warn("chmod(%s): %s", SOCK_PATH, strerror(errno));

  return fd;
}

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

  for (;;) {
    int sock = socket(AF_UNIX, SOCK_STREAM, 0);
    if (sock < 0) {
      perror("piss: socket");
      return 1;
    }

    if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) == 0) {
      int ret = 0;
      if (combo && cmd) {
        char buf[MAX_CMD_LEN * 2];
        int n = snprintf(buf, sizeof(buf), "%s|||%s", combo, cmd);
        if (n < 0 || (size_t)n >= sizeof(buf)) {
          fprintf(stderr, "piss: bind message too long\n");
          ret = 1;
        } else if (write(sock, buf, (size_t)n) < 0) {
          perror("piss: write");
          ret = 1;
        } else {
          char ack[8] = {0};
          ssize_t nr = read(sock, ack, sizeof(ack) - 1);
          if (nr > 0 && strncmp(ack, "err", 3) == 0) {
            fprintf(stderr, "piss: daemon reported bind error\n");
            ret = 1;
          }
        }
      }
      close(sock);
      return ret;
    }

    close(sock);

    if (!retry) {
      fprintf(stderr, "piss: cannot connect to %s: %s\n", SOCK_PATH,
              strerror(errno));
      return 1;
    }
    usleep(50000);
  }
}

static void get_default_config(char *buf, size_t len) {
  struct passwd *pw = getpwuid(getuid());
  if (!pw)
    die("getpwuid: %s", strerror(errno));
  int n = snprintf(buf, len, "%s/.config/piss/pissrc", pw->pw_dir);
  if (n < 0 || (size_t)n >= len)
    die("default config path too long");
}

static void ensure_executable(const char *path) {
  struct stat st;
  if (stat(path, &st) < 0)
    die("stat(%s): %s\n"
        "  (does the config file exist?)",
        path, strerror(errno));
  if (!(st.st_mode & S_IXUSR)) {
    if (chmod(path, st.st_mode | S_IXUSR) < 0)
      die("chmod(%s): %s", path, strerror(errno));
  }
}

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]);

  if (argc >= 2 &&
      (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0)) {
    fprintf(stderr, "usage:\n"
                    "  piss [pissrc]              start daemon\n"
                    "  piss wait                  wait until daemon is ready\n"
                    "  piss bind \"Mod+Key\" \"cmd\"  register a binding\n");
    return 0;
  }

  {
    struct sigaction sa;
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = SIG_IGN;
    sa.sa_flags = SA_NOCLDWAIT;
    if (sigaction(SIGCHLD, &sa, NULL) < 0)
      die("sigaction(SIGCHLD): %s", strerror(errno));
  }

  {
    struct sigaction sa;
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = handle_signal;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = 0;
    if (sigaction(SIGTERM, &sa, NULL) < 0)
      die("sigaction(SIGTERM): %s", strerror(errno));
    if (sigaction(SIGINT, &sa, NULL) < 0)
      die("sigaction(SIGINT): %s", strerror(errno));
  }

  int screen_num = 0;
  conn = xcb_connect(NULL, &screen_num);
  if (!conn)
    die("xcb_connect returned NULL — is DISPLAY set?");
  {
    int cerr = xcb_connection_has_error(conn);
    if (cerr)
      die("XCB connection error %d — is DISPLAY set correctly?", cerr);
  }

  {
    const xcb_setup_t *setup = xcb_get_setup(conn);
    xcb_screen_iterator_t it = xcb_setup_roots_iterator(setup);

    kbmap_min = setup->min_keycode;
    kbmap_max = setup->max_keycode;

    for (int i = 0; i < screen_num; i++) {
      if (!it.rem)
        die("screen %d does not exist", screen_num);
      xcb_screen_next(&it);
    }
    if (!it.rem)
      die("no screen found");
    root = it.data->root;
  }

  refresh_keymap();

  {
    uint32_t evmask = XCB_EVENT_MASK_KEY_PRESS;
    xcb_void_cookie_t ck = xcb_change_window_attributes_checked(
        conn, root, XCB_CW_EVENT_MASK, &evmask);
    xcb_generic_error_t *e = xcb_request_check(conn, ck);
    if (e) {
      int code = e->error_code;
      free(e);
      die("ChangeWindowAttributes error %d — "
          "another hotkey daemon or WM may already own the root",
          code);
    }
  }

  int serv = make_server_socket();

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

  char cfg[512];
  if (argc > 1) {
    strncpy(cfg, argv[1], sizeof(cfg) - 1);
    cfg[sizeof(cfg) - 1] = '\0';
  } else {
    get_default_config(cfg, sizeof(cfg));
  }
  ensure_executable(cfg);

  xcb_flush(conn);

  {
    pid_t pid = fork();
    if (pid < 0)
      die("fork: %s", strerror(errno));
    if (pid == 0) {
      char *args[] = {cfg, NULL};
      execvp(cfg, args);
      fprintf(stderr, "piss: execvp(%s): %s\n", cfg, strerror(errno));
      _exit(1);
    }
  }

  int xcb_fd = xcb_get_file_descriptor(conn);

  while (!g_quit) {
    while (waitpid(-1, NULL, WNOHANG) > 0)
      ;

    fd_set rfds;
    FD_ZERO(&rfds);
    FD_SET(xcb_fd, &rfds);
    FD_SET(serv, &rfds);
    int maxfd = (xcb_fd > serv) ? xcb_fd : serv;

    int r = select(maxfd + 1, &rfds, NULL, NULL, NULL);
    if (r < 0) {
      if (errno == EINTR)
        continue;
      warn("select: %s", strerror(errno));
      continue;
    }

    if (FD_ISSET(serv, &rfds)) {
      int cli = accept(serv, NULL, NULL);
      if (cli < 0) {
        warn("accept: %s", strerror(errno));
      } else {
        if (fcntl(cli, F_SETFD, FD_CLOEXEC) < 0)
          warn("fcntl(CLOEXEC) on client fd: %s", strerror(errno));

        char buf[MAX_CMD_LEN * 2];
        memset(buf, 0, sizeof(buf));
        ssize_t nr = read(cli, buf, sizeof(buf) - 1);
        if (nr > 0) {
          buf[nr] = '\0';
          int ok = add_binding(buf);
          const char *ack = (ok == 0) ? "ok\n" : "err\n";
          (void)write(cli, ack, strlen(ack));
        } else if (nr < 0) {
          warn("read(client): %s", strerror(errno));
        }
        close(cli);
      }
    }

    if (FD_ISSET(xcb_fd, &rfds)) {
      if (xcb_connection_has_error(conn)) {
        warn("XCB connection error — shutting down");
        break;
      }

      xcb_generic_event_t *ev;
      while ((ev = xcb_poll_for_event(conn)) != NULL) {
        uint8_t type = ev->response_type & ~0x80;

        switch (type) {
        case XCB_KEY_PRESS: {
          xcb_key_press_event_t *kp = (xcb_key_press_event_t *)ev;
          uint16_t state = (uint16_t)(kp->state & CLEAN_MASK);
          for (int i = 0; i < bind_count; i++) {
            if (bindings[i].keycode == kp->detail && bindings[i].mod == state) {
              spawn(bindings[i].cmd);
              break;
            }
          }
          break;
        }
        case XCB_MAPPING_NOTIFY: {
          xcb_mapping_notify_event_t *mn = (xcb_mapping_notify_event_t *)ev;
          if (mn->request == XCB_MAPPING_KEYBOARD) {
            fprintf(stderr, "piss: keyboard mapping changed, "
                            "refreshing\n");
            refresh_keymap();
          }
          break;
        }
        case 0: {
          xcb_generic_error_t *xe = (xcb_generic_error_t *)ev;
          warn("XCB error: code=%u major=%u minor=%u", xe->error_code,
               xe->major_code, xe->minor_code);
          break;
        }
        default:
          break;
        }

        free(ev);
      }
    }
  }

  ungrab_all();
  free(kbmap);
  xcb_disconnect(conn);
  unlink(SOCK_PATH);
  close(serv);

  return 0;
}