/* gcc -Wall -Wextra -ggdb -g3 -O0 -lxcb -lxcb-xkb -lxkbcommon */

#include <err.h>
#include <errno.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <xcb/xcb.h>
#include <xcb/xkb.h>
#include <xkbcommon/xkbcommon.h>

static inline bool
streq(const char *s1, const char *s2)
{
    return strcmp(s1, s2) == 0;
}

#define ARRAY_SIZE(arr) (sizeof(arr) / sizeof(*arr))

#define _cleanup_free_ __attribute__((cleanup(cleanup_free)))
static void cleanup_free(void *p) { free(*(void **)p); }

struct app {
    xcb_connection_t *conn;

    struct xkb_context *xkb_context;
    struct xkb_keymap *xkb_keymap;

    xkb_mod_index_t numlock;
    xkb_mod_index_t alt;
    xkb_mod_index_t altgr;
    xkb_mod_index_t meta;
    xkb_mod_index_t super;
    xkb_mod_index_t hyper;
};

struct vmod_masks {
    uint16_t numlock;
    uint16_t alt;
    uint16_t altgr;
    uint16_t meta;
    uint16_t super;
    uint16_t hyper;
};

static int
init_xkb(struct app *app)
{
    {
        const xcb_query_extension_reply_t *reply;
        reply = xcb_get_extension_data(app->conn, &xcb_xkb_id);
        if (!reply) {
            warnx("no XKB extension on the server");
            return -ENOTSUP;
        }
    }

    {
        xcb_xkb_use_extension_cookie_t cookie;
        _cleanup_free_ xcb_xkb_use_extension_reply_t *reply = NULL;

        cookie = xcb_xkb_use_extension(app->conn,
                                       XCB_XKB_MAJOR_VERSION,
                                       XCB_XKB_MINOR_VERSION);
        reply = xcb_xkb_use_extension_reply(app->conn, cookie, NULL);
        if (!reply) {
            warnx("couldn't start using XKB extension");
            return -EFAULT;
        }

        if (!reply->supported) {
            warnx("XKB extension version on the server is too old (want %d %d, has %d %d)",
                  XCB_XKB_MAJOR_VERSION, XCB_XKB_MINOR_VERSION,
                  reply->serverMajor, reply->serverMinor);
            return -ENOTSUP;
        }
    }

    return 0;
}

static int
app_new(xcb_connection_t *conn, struct app **out)
{
    int ret;
    struct app *app;
    static const struct xkb_rule_names names;

    app = calloc(1, sizeof(*app));
    if (!app)
        return -ENOMEM;

    app->conn = conn;

    ret = init_xkb(app);
    if (ret) {
        free(app);
        return ret;
    }

    app->xkb_context = xkb_context_new(0);
    if (!app->xkb_context) {
        warnx("couldn't get a new xkb context");
        free(app);
        return -ENOMEM;
    }

    app->xkb_keymap = xkb_keymap_new_from_names(app->xkb_context, &names, 0);
    if (!app->xkb_keymap) {
        warnx("couldn't get a new xkb context");
        xkb_context_unref(app->xkb_context);
        free(app);
        return -EFAULT;
    }

    *out = app;
    return 0;
}

static void
app_free(struct app *app)
{
    xkb_keymap_unref(app->xkb_keymap);
    xkb_context_unref(app->xkb_context);
    free(app);
}

static char *
atom_to_string(struct app *app, xcb_atom_t atom)
{
    xcb_get_atom_name_cookie_t cookie;
    _cleanup_free_ xcb_get_atom_name_reply_t *reply = NULL;
    char *pascal_str;
    int len;

    if (atom == XCB_ATOM_NONE)
        return NULL;

    cookie = xcb_get_atom_name(app->conn, atom);

    reply = xcb_get_atom_name_reply(app->conn, cookie, NULL);
    if (!reply)
        return NULL;

    pascal_str = xcb_get_atom_name_name(reply);
    len = xcb_get_atom_name_name_length(reply);

    return strndup(pascal_str, len);
}

static int
get_vmod_masks(struct app *app, struct vmod_masks *vmod_masks)
{
    xcb_xkb_get_names_cookie_t cookie;
    _cleanup_free_ xcb_xkb_get_names_reply_t *reply = NULL;
    const void *buffer;
    xcb_xkb_get_names_value_list_t value_list;
    int count;
    uint16_t vmod_mask, bit;

    cookie = xcb_xkb_get_names(app->conn,
                               XCB_XKB_ID_USE_CORE_KBD,
                               XCB_XKB_NAME_DETAIL_VIRTUAL_MOD_NAMES);

    reply = xcb_xkb_get_names_reply(app->conn, cookie, NULL);
    if (!reply) {
        warnx("couldn't get virtual modifier names from xkb");
        return -EFAULT;
    }

    buffer = xcb_xkb_get_names_value_list(reply);
    xcb_xkb_get_names_value_list_unpack(buffer,
                                        reply->nTypes,
                                        reply->indicators,
                                        reply->virtualMods,
                                        reply->groupNames,
                                        reply->nKeys,
                                        reply->nKeyAliases,
                                        reply->nRadioGroups,
                                        reply->which,
                                        &value_list);

    vmod_mask = reply->virtualMods;
    count = 0;
    for (bit = 1; vmod_mask; bit <<= 1) {
        _cleanup_free_ char *vmod_name = NULL;

        if (!(vmod_mask & bit))
            continue;

        vmod_mask &= ~bit;
        vmod_name = atom_to_string(app, value_list.virtualModNames[count]);
        count++;

        if (!vmod_name)
            continue;

        if (streq(vmod_name, "NumLock"))
            vmod_masks->numlock = bit;
        else if (streq(vmod_name, "Alt"))
            vmod_masks->alt = bit;
        else if (streq(vmod_name, "AltGr"))
            vmod_masks->altgr = bit;
        else if (streq(vmod_name, "Meta"))
            vmod_masks->meta = bit;
        else if (streq(vmod_name, "Super"))
            vmod_masks->super = bit;
        else if (streq(vmod_name, "Hyper"))
            vmod_masks->hyper = bit;
    }

    return 0;
}

static xkb_mod_index_t
modmap_to_xkb_mod_index(struct app *app, uint8_t modmap)
{
    static const struct {
        xcb_mod_mask_t mask;
        const char *name;
    } mask_to_name[] = {
        { XCB_MOD_MASK_SHIFT, "Shift" },
        { XCB_MOD_MASK_LOCK, "Lock" },
        { XCB_MOD_MASK_CONTROL, "Control" },
        { XCB_MOD_MASK_1, "Mod1" },
        { XCB_MOD_MASK_2, "Mod2" },
        { XCB_MOD_MASK_3, "Mod3" },
        { XCB_MOD_MASK_4, "Mod4" },
        { XCB_MOD_MASK_5, "Mod5" },
    };
    unsigned int i;
    struct xkb_keymap *keymap = app->xkb_keymap;

    for (i = 0; i < ARRAY_SIZE(mask_to_name); i++)
        if (modmap & mask_to_name[i].mask)
            return xkb_keymap_mod_get_index(keymap, mask_to_name[i].name);

    return XKB_MOD_INVALID;
}

static int
get_vmod_mappings(struct app *app, const struct vmod_masks *vmod_masks)
{
    xcb_xkb_get_map_cookie_t cookie;
    _cleanup_free_ xcb_xkb_get_map_reply_t *reply = NULL;
    xcb_xkb_get_map_map_t map;
    void *buffer;
    int count;
    uint16_t vmod_mask, bit;

    cookie = xcb_xkb_get_map(app->conn,
                             XCB_XKB_ID_USE_CORE_KBD,
                             XCB_XKB_MAP_PART_VIRTUAL_MODS,
                             0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);

    reply = xcb_xkb_get_map_reply(app->conn, cookie, NULL);
    if (!reply) {
        warnx("couldn't get xkb map");
        return -EFAULT;
    }

    buffer = xcb_xkb_get_map_map(reply);
    xcb_xkb_get_map_map_unpack(buffer,
                               reply->nTypes,
                               reply->nKeySyms,
                               reply->nKeyActions,
                               reply->totalActions,
                               reply->totalKeyBehaviors,
                               reply->nVModMapKeys,
                               reply->totalKeyExplicit,
                               reply->totalModMapKeys,
                               reply->totalVModMapKeys,
                               reply->present,
                               &map);

    vmod_mask = reply->virtualMods;
    count = 0;
    for (bit = 1; vmod_mask; bit <<= 1) {
        uint8_t modmap;
        xkb_mod_index_t mod_index;

        if (!(vmod_mask & bit))
            continue;

        vmod_mask &= ~bit;
        modmap = map.vmods_rtrn[count];
        count++;

        mod_index = modmap_to_xkb_mod_index(app, modmap);
        if (mod_index == XKB_MOD_INVALID)
            continue;

        if (vmod_masks->numlock == bit)
            app->numlock = mod_index;
        else if (vmod_masks->alt == bit)
            app->alt = mod_index;
        else if (vmod_masks->altgr == bit)
            app->altgr = mod_index;
        else if (vmod_masks->meta == bit)
            app->meta = mod_index;
        else if (vmod_masks->super == bit)
            app->super = mod_index;
        else if (vmod_masks->hyper == bit)
            app->hyper = mod_index;
    }

    return 0;
}

static int get_vmod_to_real_mod_mappings(struct app *app)
{
    int ret;
    struct vmod_masks vmod_masks = { 0x0000 };

    ret = get_vmod_masks(app, &vmod_masks);
    if (ret)
        return ret;

    ret = get_vmod_mappings(app, &vmod_masks);
    if (ret)
        return ret;

    return 0;
}

#define _cleanup_app_ __attribute__((cleanup(cleanup_app)))
static void cleanup_app(struct app **app) { app_free(*app); }
#define _cleanup_conn_ __attribute__((cleanup(cleanup_conn)))
static void cleanup_conn(xcb_connection_t **conn) { if (*conn) xcb_disconnect(*conn); }

int
main(void)
{
    int ret;
    _cleanup_app_ struct app *app = NULL;
    _cleanup_conn_ xcb_connection_t *conn = NULL;
    struct xkb_keymap *keymap;

    conn = xcb_connect(NULL, NULL);
    if (xcb_connection_has_error(conn)) {
        warnx("couldn't connect to display");
        return EXIT_FAILURE;
    }

    ret = app_new(conn, &app);
    if (ret) {
        warnx("couldn't initialize app");
        return EXIT_FAILURE;
    }

    ret = get_vmod_to_real_mod_mappings(app);
    if (ret)
        return EXIT_FAILURE;

    keymap = app->xkb_keymap;
    printf("NumLock is %s\n", xkb_keymap_mod_get_name(keymap, app->numlock));
    printf("Alt is %s\n", xkb_keymap_mod_get_name(keymap, app->alt));
    printf("AltGr is %s\n", xkb_keymap_mod_get_name(keymap, app->altgr));
    printf("Meta is %s\n", xkb_keymap_mod_get_name(keymap, app->meta));
    printf("Super is %s\n", xkb_keymap_mod_get_name(keymap, app->super));
    printf("Hyper is %s\n", xkb_keymap_mod_get_name(keymap, app->hyper));

    return EXIT_SUCCESS;
}