#include <X11/XKBlib.h>
#include <X11/Xlib.h>
#include <X11/keysym.h>
#include <pwd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.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;
  KeySym keysym;
  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 *mod_str) {
  unsigned int mod = 0;
  char *str = strdup(mod_str);
  char *token = strtok(str, "+");

  while (token) {
    while (*token == ' ')
      token++;
    char *end = token + strlen(token) - 1;
    while (end > token && *end == ' ')
      end--;
    *(end + 1) = '\0';

    if (strcasecmp(token, "super") == 0 || strcasecmp(token, "mod4") == 0)
      mod |= Mod4Mask;
    else if (strcasecmp(token, "shift") == 0)
      mod |= ShiftMask;
    else if (strcasecmp(token, "ctrl") == 0 ||
             strcasecmp(token, "control") == 0)
      mod |= ControlMask;
    else if (strcasecmp(token, "alt") == 0 || strcasecmp(token, "mod1") == 0)
      mod |= Mod1Mask;

    token = strtok(NULL, "+");
  }

  free(str);
  return mod;
}

KeySym parse_key(const char *key_str) {
  char *str = strdup(key_str);
  char *s = str;
  while (*s == ' ')
    s++;
  char *end = s + strlen(s) - 1;
  while (end > s && *end == ' ')
    end--;
  *(end + 1) = '\0';

  KeySym ks = XStringToKeysym(s);
  free(str);
  return ks;
}

void add_binding(const char *keycombo, const char *cmd) {
  if (bind_count >= MAX_BINDS) {
    fprintf(stderr, "piss: max bindings reached\n");
    return;
  }

  char *combo = strdup(keycombo);
  char *end = combo + strlen(combo) - 1;
  while (end > combo && (*end == ' ' || *end == '\t' || *end == '\n')) {
    *end = '\0';
    end--;
  }

  char *plus = strrchr(combo, '+');

  if (!plus) {
    fprintf(stderr, "piss: invalid key combo: %s\n", keycombo);
    free(combo);
    return;
  }

  *plus = '\0';
  char *mod_str = combo;
  char *key_str = plus + 1;

  bindings[bind_count].mod = parse_modifiers(mod_str);
  bindings[bind_count].keysym = parse_key(key_str);

  if (bindings[bind_count].keysym == NoSymbol) {
    fprintf(stderr, "piss: invalid key: %s (from %s)\n", key_str, keycombo);
    free(combo);
    return;
  }

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

  KeyCode kc = XKeysymToKeycode(dpy, bindings[bind_count].keysym);
  if (kc != 0) {
    XSync(dpy, False);

    XGrabKey(dpy, kc, bindings[bind_count].mod, root, True, GrabModeAsync,
             GrabModeAsync);
    XGrabKey(dpy, kc, bindings[bind_count].mod | Mod2Mask, root, True,
             GrabModeAsync, GrabModeAsync);
    XGrabKey(dpy, kc, bindings[bind_count].mod | LockMask, root, True,
             GrabModeAsync, GrabModeAsync);
    XGrabKey(dpy, kc, bindings[bind_count].mod | Mod2Mask | LockMask, root,
             True, GrabModeAsync, GrabModeAsync);

    XSync(dpy, False);

    fprintf(stderr, "piss: bound [%s] keysym=%lu keycode=%u mod=%u -> %s\n",
            keycombo, bindings[bind_count].keysym, kc, bindings[bind_count].mod,
            cmd);

    bind_count++;
  } else {
    fprintf(stderr, "piss: couldn't get keycode for: %s\n", keycombo);
  }

  free(combo);
}

void execute_command(const char *cmd) {
  if (fork() == 0) {
    setsid();
    execl("/bin/sh", "sh", "-c", cmd, NULL);
    exit(0);
  }
}

void handle_keypress(XEvent *ev) {
  XKeyEvent *e = &ev->xkey;
  KeySym ks = XkbKeycodeToKeysym(dpy, e->keycode, 0, 0);
  unsigned int cleanmask = e->state & ~(Mod2Mask | LockMask);

  fprintf(stderr, "piss: key pressed: keysym=%lu mod=%u\n", ks, cleanmask);

  for (int i = 0; i < bind_count; i++) {
    if (bindings[i].keysym == ks && bindings[i].mod == cleanmask) {
      fprintf(stderr, "piss: executing: %s\n", bindings[i].cmd);
      execute_command(bindings[i].cmd);
      break;
    }
  }
}

char *get_default_config() {
  static char path[512];
  const char *home = getenv("HOME");
  if (!home) {
    struct passwd *pw = getpwuid(getuid());
    home = pw->pw_dir;
  }
  snprintf(path, sizeof(path), "%s/.config/piss/pissrc", home);
  return path;
}

int send_bind_to_daemon(const char *keycombo, const char *cmd) {
  int sock;
  struct sockaddr_un addr;
  char msg[MAX_CMD_LEN * 2];

  sock = socket(AF_UNIX, SOCK_STREAM, 0);
  if (sock < 0) {
    return 1;
  }

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

  if (connect(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
    close(sock);
    return 1;
  }

  snprintf(msg, sizeof(msg), "%s|||%s", keycombo, cmd);
  write(sock, msg, strlen(msg));
  close(sock);
  return 0;
}

void handle_socket_connection(int client_sock) {
  char buf[MAX_CMD_LEN * 2];
  int n = read(client_sock, buf, sizeof(buf) - 1);
  if (n > 0) {
    buf[n] = '\0';

    char *sep = strstr(buf, "|||");
    if (sep) {
      *sep = '\0';
      char *keycombo = buf;
      char *cmd = sep + 3;
      add_binding(keycombo, cmd);
    }
  }
  close(client_sock);
}

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

  char *config = (argc >= 2) ? argv[1] : get_default_config();

  dpy = XOpenDisplay(NULL);
  if (!dpy) {
    fprintf(stderr, "piss: cannot open display\n");
    return 1;
  }

  root = DefaultRootWindow(dpy);

  int sock, client_sock;
  struct sockaddr_un addr;

  unlink(SOCK_PATH);
  sock = socket(AF_UNIX, SOCK_STREAM, 0);
  if (sock < 0) {
    perror("socket");
    return 1;
  }

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

  if (bind(sock, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
    perror("bind");
    return 1;
  }

  if (listen(sock, 5) < 0) {
    perror("listen");
    return 1;
  }

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

  pid_t config_pid = fork();
  if (config_pid == 0) {
    usleep(50000);

    FILE *f = fopen(config, "r");
    if (f) {
      char line[256];
      if (fgets(line, sizeof(line), f) && strncmp(line, "#!", 2) == 0) {
        fclose(f);
        execl(config, config, NULL);
      } else {
        fclose(f);
        execl("/bin/sh", "sh", config, NULL);
      }
    }

    perror("execl");
    exit(1);
  }

  fd_set readfds;
  struct timeval tv;
  int status;

  while (waitpid(config_pid, &status, WNOHANG) == 0) {
    FD_ZERO(&readfds);
    FD_SET(sock, &readfds);
    tv.tv_sec = 0;
    tv.tv_usec = 50000;

    if (select(sock + 1, &readfds, NULL, NULL, &tv) > 0) {
      client_sock = accept(sock, NULL, NULL);
      if (client_sock >= 0) {
        handle_socket_connection(client_sock);
      }
    }
  }

  for (int i = 0; i < 5; i++) {
    FD_ZERO(&readfds);
    FD_SET(sock, &readfds);
    tv.tv_sec = 0;
    tv.tv_usec = 10000;

    if (select(sock + 1, &readfds, NULL, NULL, &tv) > 0) {
      client_sock = accept(sock, NULL, NULL);
      if (client_sock >= 0) {
        handle_socket_connection(client_sock);
      }
    }
  }

  XSelectInput(dpy, root, KeyPressMask);

  printf("piss: running with %d bindings\n", bind_count);

  int x11_fd = ConnectionNumber(dpy);

  fprintf(stderr, "piss: listening for keypresses (X11 fd=%d)...\n", x11_fd);

  while (1) {
    FD_ZERO(&readfds);
    FD_SET(x11_fd, &readfds);
    FD_SET(sock, &readfds);
    int max_fd = (x11_fd > sock) ? x11_fd : sock;

    if (select(max_fd + 1, &readfds, NULL, NULL, NULL) > 0) {
      if (FD_ISSET(x11_fd, &readfds)) {
        while (XPending(dpy)) {
          XEvent ev;
          XNextEvent(dpy, &ev);
          fprintf(stderr, "piss: received X event type=%d\n", ev.type);
          if (ev.type == KeyPress)
            handle_keypress(&ev);
        }
      }

      if (FD_ISSET(sock, &readfds)) {
        client_sock = accept(sock, NULL, NULL);
        if (client_sock >= 0) {
          handle_socket_connection(client_sock);
        }
      }
    }
  }

  unlink(SOCK_PATH);
  XCloseDisplay(dpy);
  return 0;
}