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

#include <errno.h>
#include <fcntl.h>
#include <limits.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/chd.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;
static xcb_window_t root;

static xcb_get_keyboard_mapping_reply_t *kbmap;
static xcb_keycode_t kbmap_min;
static xcb_keycode_t kbmap_max;

static Binding bindings[MAX_BINDS];
static int bind_count;

static volatile sig_atomic_t g_quit;
static volatile sig_atomic_t g_reload;
static volatile sig_atomic_t g_restart;

static char g_cfg[512];
static char **g_argv;

static void die(const char *msg) {
    fprintf(stderr, "chd: %s\n", msg);
    exit(1);
}

static void die_errno(const char *msg) {
    fprintf(stderr, "chd: %s: %s\n", msg, strerror(errno));
    exit(1);
}

static void handle_signal(int sig) {
    if (sig == SIGTERM || sig == SIGINT)
        g_quit = 1;
    if (sig == SIGUSR1)
        g_reload = 1;
    if (sig == SIGUSR2)
        g_restart = 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) {
        fprintf(stderr, "chd: keyboard mapping failed (error %d)\n", e->error_code);
        free(e);
    }
    if (!kbmap)
        die("could not retrieve keyboard mapping");
}

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 i = 0; i < kpk; i++)
            if (row[i] == 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),
    };

    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) {
            fprintf(stderr, "chd: grab_key failed for mod=%u kc=%u (error %d)\n",
                    mods[i], kc, e->error_code);
            free(e);
            return -1;
        }
    }
    return 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 void spawn(const char *cmd) {
    pid_t pid = fork();
    if (pid < 0) {
        fprintf(stderr, "chd: fork failed: %s\n", strerror(errno));
        return;
    }
    if (pid == 0) {
        setsid();
        execl("/bin/sh", "sh", "-c", cmd, (char *)NULL);
        fprintf(stderr, "chd: execl failed: %s\n", strerror(errno));
        _exit(1);
    }
}

static xcb_window_t find_first_window(xcb_window_t w) {
    xcb_query_tree_cookie_t ck = xcb_query_tree(conn, w);
    xcb_generic_error_t *e = NULL;
    xcb_query_tree_reply_t *r = xcb_query_tree_reply(conn, ck, &e);
    if (e) {
        free(e);
        return XCB_NONE;
    }
    if (!r)
        return XCB_NONE;

    xcb_window_t *children = xcb_query_tree_children(r);
    int n = xcb_query_tree_children_length(r);
    xcb_window_t found = XCB_NONE;

    for (int i = 0; i < n && found == XCB_NONE; i++) {
        xcb_get_window_attributes_cookie_t ack =
            xcb_get_window_attributes(conn, children[i]);
        xcb_get_window_attributes_reply_t *attr =
            xcb_get_window_attributes_reply(conn, ack, NULL);

        if (attr) {
            if (attr->map_state == XCB_MAP_STATE_VIEWABLE &&
                attr->_class == XCB_WINDOW_CLASS_INPUT_OUTPUT) {
                found = children[i];
            }
            free(attr);
        }

        if (found == XCB_NONE)
            found = find_first_window(children[i]);
    }

    free(r);
    return found;
}

static int focus_first_window(void) {
    xcb_window_t w = find_first_window(root);
    if (w == XCB_NONE) {
        fprintf(stderr, "chd: no viewable window found\n");
        return 1;
    }

    xcb_void_cookie_t ck = xcb_set_input_focus_checked(
        conn, XCB_INPUT_FOCUS_POINTER_ROOT, w, XCB_CURRENT_TIME);
    xcb_generic_error_t *e = xcb_request_check(conn, ck);
    if (e) {
        fprintf(stderr, "chd: set_input_focus failed (error %d)\n", e->error_code);
        free(e);
        return 1;
    }

    xcb_flush(conn);
    return 0;
}

static int add_binding(char *msg) {
    if (!strncmp(msg, "reload", 6)) {
        g_reload = 1;
        return 0;
    }

    if (!strncmp(msg, "restart", 7)) {
        g_restart = 1;
        return 0;
    }

    if (!strncmp(msg, "focus", 5))
        return focus_first_window();

    char *sep = strstr(msg, "|||");
    if (!sep) {
        fprintf(stderr, "chd: malformed bind message (missing |||)\n");
        return -1;
    }

    *sep = 0;
    char *combo = msg;
    char *cmd = sep + 3;

    char *last = strrchr(combo, '+');
    if (!last) {
        fprintf(stderr, "chd: malformed combo (no + separator)\n");
        return -1;
    }

    *last = 0;
    char *key = last + 1;

    KeySym sym = XStringToKeysym(key);
    if (sym == NoSymbol) {
        fprintf(stderr, "chd: unknown key '%s'\n", key);
        return -1;
    }

    xcb_keycode_t kc = keysym_to_keycode(sym);
    if (!kc) {
        fprintf(stderr, "chd: no keycode for key '%s'\n", key);
        return -1;
    }

    uint16_t mod = parse_modifiers(combo);

    for (int i = 0; i < bind_count; i++) {
        if (bindings[i].mod == mod && bindings[i].keycode == kc) {
            strncpy(bindings[i].cmd, cmd, MAX_CMD_LEN - 1);
            bindings[i].cmd[MAX_CMD_LEN - 1] = '\0';
            return 0;
        }
    }

    if (bind_count >= MAX_BINDS) {
        fprintf(stderr, "chd: binding table full (%d max)\n", MAX_BINDS);
        return -1;
    }

    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';

    if (grab_key(mod, kc) < 0)
        return -1;

    bind_count++;
    return 0;
}

static int make_socket(void) {
    int fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0)
        die_errno("socket");

    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_errno("bind");

    if (listen(fd, 16) < 0)
        die_errno("listen");

    return fd;
}

static int client(const char *msg) {
    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);

    int s = socket(AF_UNIX, SOCK_STREAM, 0);
    if (s < 0) {
        fprintf(stderr, "chd: socket: %s\n", strerror(errno));
        return 1;
    }

    if (connect(s, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        fprintf(stderr, "chd: connect: %s\n", strerror(errno));
        close(s);
        return 1;
    }

    ssize_t w = write(s, msg, strlen(msg));
    if (w < 0)
        fprintf(stderr, "chd: write: %s\n", strerror(errno));

    close(s);
    return w < 0 ? 1 : 0;
}

static void get_cfg(void) {
    struct passwd *pw = getpwuid(getuid());
    if (!pw)
        die_errno("getpwuid");
    snprintf(g_cfg, sizeof(g_cfg), "%s/.config/chd/chdrc", pw->pw_dir);
}

int main(int argc, char *argv[]) {
    g_argv = argv;

    if (argc > 1 && !strcmp(argv[1], "bind")) {
        if (argc < 3) {
            fprintf(stderr, "chd: bind requires an argument\n");
            return 1;
        }
        return client(argv[2]);
    }

    if (argc > 1 && !strcmp(argv[1], "reload"))
        return client("reload");

    if (argc > 1 && !strcmp(argv[1], "restart"))
        return client("restart");

    if (argc > 1 && !strcmp(argv[1], "focus")) {
        conn = xcb_connect(NULL, NULL);
        if (xcb_connection_has_error(conn)) {
            fprintf(stderr, "chd: xcb_connect failed\n");
            return 1;
        }
        root = xcb_setup_roots_iterator(xcb_get_setup(conn)).data->root;
        int ret = focus_first_window();
        xcb_disconnect(conn);
        return ret;
    }

    if (argc > 1) {
        fprintf(stderr, "chd: unknown argument '%s'\n", argv[1]);
        return 1;
    }

    struct sigaction sa;
    memset(&sa, 0, sizeof(sa));
    sa.sa_handler = handle_signal;
    sigaction(SIGINT, &sa, NULL);
    sigaction(SIGTERM, &sa, NULL);
    sigaction(SIGUSR1, &sa, NULL);
    sigaction(SIGUSR2, &sa, NULL);

    conn = xcb_connect(NULL, NULL);
    if (xcb_connection_has_error(conn))
        die("xcb_connect failed");

    const xcb_setup_t *setup = xcb_get_setup(conn);
    xcb_screen_t *screen = xcb_setup_roots_iterator(setup).data;
    root = screen->root;
    kbmap_min = setup->min_keycode;
    kbmap_max = setup->max_keycode;

    refresh_keymap();
    get_cfg();

    int serv = make_socket();
    int xfd = xcb_get_file_descriptor(conn);

    while (!g_quit) {
        if (xcb_connection_has_error(conn)) {
            fprintf(stderr, "chd: lost xcb connection\n");
            break;
        }

        if (g_reload) {
            g_reload = 0;
            ungrab_all();
            bind_count = 0;
            refresh_keymap();
        }

        if (g_restart) {
            ungrab_all();
            xcb_disconnect(conn);
            close(serv);
            execvp(argv[0], argv);
            die_errno("execvp");
        }

        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(xfd, &fds);
        FD_SET(serv, &fds);
        int nfds = (xfd > serv ? xfd : serv) + 1;

        int r = select(nfds, &fds, NULL, NULL, NULL);
        if (r < 0) {
            if (errno == EINTR)
                continue;
            fprintf(stderr, "chd: select: %s\n", strerror(errno));
            break;
        }

        if (FD_ISSET(serv, &fds)) {
            int c = accept(serv, NULL, NULL);
            if (c < 0) {
                fprintf(stderr, "chd: accept: %s\n", strerror(errno));
            } else {
                char buf[2048];
                int n = read(c, buf, sizeof(buf) - 1);
                if (n > 0) {
                    buf[n] = 0;
                    add_binding(buf);
                } else if (n < 0) {
                    fprintf(stderr, "chd: read: %s\n", strerror(errno));
                }
                close(c);
            }
        }

        if (FD_ISSET(xfd, &fds)) {
            xcb_generic_event_t *ev;
            while ((ev = xcb_poll_for_event(conn))) {
                if ((ev->response_type & 0x7f) == XCB_KEY_PRESS) {
                    xcb_key_press_event_t *kp = (xcb_key_press_event_t *)ev;
                    uint16_t st = kp->state & CLEAN_MASK;

                    for (int i = 0; i < bind_count; i++) {
                        if (bindings[i].keycode == kp->detail &&
                            bindings[i].mod == st) {
                            spawn(bindings[i].cmd);
                            break;
                        }
                    }
                }
                free(ev);
            }
        }
    }

    ungrab_all();
    xcb_disconnect(conn);
    close(serv);
    return 0;
}