/* Minimal x86-64 Linux hotkey daemon.  Intentionally uses no C library. */

typedef unsigned char u8;
typedef unsigned short u16;
typedef unsigned int u32;
typedef unsigned long u64;
typedef long i64;

#define SYS_read                 0
#define SYS_close                3
#define SYS_ioctl               16
#define SYS_rt_sigaction        13
#define SYS_dup2                33
#define SYS_fork                57
#define SYS_execve              59
#define SYS_exit                60
#define SYS_openat             257
#define SYS_inotify_add_watch  254
#define SYS_epoll_wait         232
#define SYS_epoll_ctl          233
#define SYS_epoll_create1      291
#define SYS_inotify_init1      294

#define AT_FDCWD              -100
#define O_RDONLY                 0
#define O_RDWR                   2
#define O_NONBLOCK            04000
#define O_CLOEXEC           02000000

#define EPOLLIN             0x001
#define EPOLLERR            0x008
#define EPOLLHUP            0x010
#define EPOLL_CTL_ADD           1
#define EPOLL_CTL_DEL           2
#define EPOLL_CLOEXEC      02000000

#define IN_MOVED_FROM       0x040
#define IN_MOVED_TO         0x080
#define IN_CREATE           0x100
#define IN_DELETE           0x200
#define IN_NONBLOCK        00004000
#define IN_CLOEXEC        02000000

#define EV_KEY                  1
#define KEY_LEFTALT            56
#define KEY_SYSRQ              99
#define KEY_RIGHTALT          100

#define SIGCHLD                17
#define SIG_IGN                 1

#define MAX_KEYBOARDS          16
#define MAX_EPOLL_EVENTS       16
#define MAX_EVENT_NUMBER       63
#define TAG_INOTIFY 0xffffffffffffffffUL

struct epoll_event {
    u32 events;
    u64 data;
} __attribute__((packed));

struct input_event {
    i64 seconds;
    i64 microseconds;
    u16 type;
    u16 code;
    u32 value;
};

struct kernel_sigaction {
    void (*handler)(int);
    u64 flags;
    void (*restorer)(void);
    u64 mask;
};

struct keyboard {
    int fd;
    u8 left_alt;
    u8 right_alt;
};

static struct keyboard keyboards[MAX_KEYBOARDS];
static int keyboard_count;
static int epoll_fd;
static int inotify_fd;

static i64 syscall1(i64 number, i64 a1)
{
    i64 result;
    __asm__ volatile("syscall" : "=a"(result) : "a"(number), "D"(a1)
                     : "rcx", "r11", "memory");
    return result;
}

static i64 syscall2(i64 number, i64 a1, i64 a2)
{
    i64 result;
    __asm__ volatile("syscall" : "=a"(result)
                     : "a"(number), "D"(a1), "S"(a2)
                     : "rcx", "r11", "memory");
    return result;
}

static i64 syscall3(i64 number, i64 a1, i64 a2, i64 a3)
{
    i64 result;
    __asm__ volatile("syscall" : "=a"(result)
                     : "a"(number), "D"(a1), "S"(a2), "d"(a3)
                     : "rcx", "r11", "memory");
    return result;
}

static i64 syscall4(i64 number, i64 a1, i64 a2, i64 a3, i64 a4)
{
    register i64 r10 __asm__("r10") = a4;
    i64 result;
    __asm__ volatile("syscall" : "=a"(result)
                     : "a"(number), "D"(a1), "S"(a2), "d"(a3), "r"(r10)
                     : "rcx", "r11", "memory");
    return result;
}

__attribute__((noreturn)) static void terminate(int status)
{
    syscall1(SYS_exit, status);
    __builtin_unreachable();
}

static int open_file(const char *path, int flags)
{
    return (int)syscall4(SYS_openat, AT_FDCWD, (i64)path, flags, 0);
}

static void close_file(int fd)
{
    syscall1(SYS_close, fd);
}

static void remove_keyboard(int index)
{
    syscall4(SYS_epoll_ctl, epoll_fd, EPOLL_CTL_DEL,
             keyboards[index].fd, 0);
    close_file(keyboards[index].fd);
    keyboards[index] = keyboards[--keyboard_count];
}

static int is_printscreen_device(int fd)
{
    u64 key_bits[2] = {0, 0};
    u64 request = (2UL << 30) | (sizeof(key_bits) << 16) |
                  ((u64)'E' << 8) | (0x20 + EV_KEY);
    if (syscall3(SYS_ioctl, fd, request, (i64)key_bits) < 0)
        return 0;
    return (key_bits[KEY_SYSRQ / 64] >> (KEY_SYSRQ % 64)) & 1;
}

static void make_event_path(char path[22], int number)
{
    static const char prefix[] = "/dev/input/event";
    int i;
    for (i = 0; i < 16; ++i)
        path[i] = prefix[i];
    if (number >= 10) {
        path[i++] = (char)('0' + number / 10);
        path[i++] = (char)('0' + number % 10);
    } else {
        path[i++] = (char)('0' + number);
    }
    path[i] = 0;
}

static void add_keyboard_fd(int fd)
{
    if (keyboard_count == MAX_KEYBOARDS) {
        close_file(fd);
        return;
    }
    struct epoll_event event = {
        .events = EPOLLIN | EPOLLERR | EPOLLHUP,
        .data = (u64)(u32)fd,
    };
    if (syscall4(SYS_epoll_ctl, epoll_fd, EPOLL_CTL_ADD, fd,
                 (i64)&event) < 0) {
        close_file(fd);
        return;
    }
    keyboards[keyboard_count].fd = fd;
    keyboards[keyboard_count].left_alt = 0;
    keyboards[keyboard_count].right_alt = 0;
    ++keyboard_count;
}

static void scan_keyboards(void)
{
    while (keyboard_count)
        remove_keyboard(keyboard_count - 1);

    for (int number = 0; number <= MAX_EVENT_NUMBER; ++number) {
        char path[22];
        make_event_path(path, number);
        int fd = open_file(path, O_RDONLY | O_NONBLOCK | O_CLOEXEC);
        if (fd < 0)
            continue;
        if (is_printscreen_device(fd))
            add_keyboard_fd(fd);
        else
            close_file(fd);
    }
}

static int alt_is_down(void)
{
    for (int i = 0; i < keyboard_count; ++i) {
        if (keyboards[i].left_alt || keyboards[i].right_alt)
            return 1;
    }
    return 0;
}

static void take_screenshot(int window)
{
    i64 child = syscall1(SYS_fork, 0);
    if (child != 0)
        return;

    int null_fd = open_file("/dev/null", O_RDWR);
    if (null_fd >= 0) {
        syscall2(SYS_dup2, null_fd, 0);
        syscall2(SYS_dup2, null_fd, 1);
        syscall2(SYS_dup2, null_fd, 2);
        if (null_fd > 2)
            close_file(null_fd);
    }

    static char *const environment[] = {
        "PATH=/usr/bin:/bin",
        0,
    };
    static char *const full_screen[] = {
        "/usr/bin/runuser", "-u", "cachyos", "--", "/usr/bin/env",
        "DISPLAY=:0", "XAUTHORITY=/home/cachyos/.Xauthority",
        "DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/1000/bus",
        "/usr/bin/gnome-screenshot", "--clipboard", 0,
    };
    static char *const active_window[] = {
        "/usr/bin/runuser", "-u", "cachyos", "--", "/usr/bin/env",
        "DISPLAY=:0", "XAUTHORITY=/home/cachyos/.Xauthority",
        "DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/1000/bus",
        "/usr/bin/gnome-screenshot", "--window", "--clipboard", 0,
    };
    char *const *arguments = window ? active_window : full_screen;
    syscall3(SYS_execve, (i64)arguments[0], (i64)arguments,
             (i64)environment);
    terminate(127);
}

static void handle_keyboard(int index)
{
    struct input_event events[16];
    i64 bytes;
    while ((bytes = syscall3(SYS_read, keyboards[index].fd,
                             (i64)events, sizeof(events))) > 0) {
        int count = (int)(bytes / (i64)sizeof(events[0]));
        for (int i = 0; i < count; ++i) {
            if (events[i].type != EV_KEY)
                continue;
            u8 pressed = events[i].value != 0;
            if (events[i].code == KEY_LEFTALT)
                keyboards[index].left_alt = pressed;
            else if (events[i].code == KEY_RIGHTALT)
                keyboards[index].right_alt = pressed;
            else if (events[i].code == KEY_SYSRQ && events[i].value == 1)
                take_screenshot(alt_is_down());
        }
    }
    /* Linux returns -EAGAIN (-11) after all pending nonblocking events. */
    if (bytes != -11)
        remove_keyboard(index);
}

static void drain_inotify(void)
{
    char buffer[512];
    while (syscall3(SYS_read, inotify_fd, (i64)buffer, sizeof(buffer)) > 0)
        ;
    scan_keyboards();
}

static int daemon_main(void)
{
    struct kernel_sigaction action = {
        .handler = (void (*)(int))SIG_IGN,
        .flags = 0,
        .restorer = 0,
        .mask = 0,
    };
    syscall4(SYS_rt_sigaction, SIGCHLD, (i64)&action, 0, 8);

    epoll_fd = (int)syscall1(SYS_epoll_create1, EPOLL_CLOEXEC);
    inotify_fd = (int)syscall1(SYS_inotify_init1,
                               IN_NONBLOCK | IN_CLOEXEC);
    if (epoll_fd < 0 || inotify_fd < 0)
        return 1;

    if (syscall3(SYS_inotify_add_watch, inotify_fd,
                 (i64)"/dev/input",
                 IN_CREATE | IN_DELETE | IN_MOVED_FROM | IN_MOVED_TO) < 0)
        return 1;

    struct epoll_event notify_event = {
        .events = EPOLLIN,
        .data = TAG_INOTIFY,
    };
    if (syscall4(SYS_epoll_ctl, epoll_fd, EPOLL_CTL_ADD, inotify_fd,
                 (i64)&notify_event) < 0)
        return 1;

    scan_keyboards();

    for (;;) {
        struct epoll_event events[MAX_EPOLL_EVENTS];
        i64 count = syscall4(SYS_epoll_wait, epoll_fd, (i64)events,
                             MAX_EPOLL_EVENTS, -1);
        if (count < 0)
            continue;
        for (i64 i = 0; i < count; ++i) {
            if (events[i].data == TAG_INOTIFY) {
                drain_inotify();
                continue;
            }
            int fd = (int)events[i].data;
            for (int index = 0; index < keyboard_count; ++index) {
                if (keyboards[index].fd == fd) {
                    handle_keyboard(index);
                    break;
                }
            }
        }
    }
}

__attribute__((noreturn)) void _start(void)
{
    terminate(daemon_main());
}