#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <dirent.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/connector.h>
#include <linux/cn_proc.h>

#define MAX_GAMES        32
#define ENVIRON_BUF_SIZE (64 * 1024)
#define CMDLINE_BUF_SIZE 4096
#define RECV_BUF_SIZE    8192

static pid_t games[MAX_GAMES]; /* 0 = empty slot */
static int   game_count = 0;
static char  environ_buf[ENVIRON_BUF_SIZE];
static char  cmdline_buf[CMDLINE_BUF_SIZE];
static char  recv_buf[RECV_BUF_SIZE];

static int find_slot_by_pid(pid_t pid) {
    for (int i = 0; i < MAX_GAMES; i++) {
        if (games[i] == pid)
            return i;
    }
    return -1;
}

static int find_free_slot(void) {
    for (int i = 0; i < MAX_GAMES; i++) {
        if (!games[i])
            return i;
    }
    return -1;
}

static void add_game(pid_t pid) {
    int slot = find_free_slot();
    if (slot < 0) return;
    games[slot] = pid;
    if (game_count++ == 0) {
        printf("Game launched\n");
        fflush(stdout);
    }
}

static void remove_game(pid_t pid) {
    int slot = find_slot_by_pid(pid);
    if (slot < 0) return;
    games[slot] = 0;
    if (--game_count == 0) {
        printf("Game exited\n");
        fflush(stdout);
    }
}

static ssize_t read_file(const char *path, char *buf, size_t size) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) return -1;
    ssize_t n = read(fd, buf, size - 1);
    close(fd);
    if (n > 0) buf[n] = '\0';
    return n;
}

/* Return pointer to value inside null-separated environ block, or NULL. */
static const char *env_get(const char *env, size_t size, const char *key) {
    const size_t klen = strlen(key);
    size_t index = 0;
    while (index < size) {
        const char *env_start = env + index;
        const char *env_end = memchr(env_start, '\0', size - index);
        if(!env_end)
            break;

        const size_t env_len = env_end - env_start;
        if(env_len >= klen + 1 && memcmp(env_start, key, klen) == 0 && env_start[klen] == '=')
            return env_start + klen + 1;

        index += env_len + 1;
    }
    return NULL;
}

static size_t get_argv_len(const char *cmdline, size_t size) {
    const char *argv0_end = memchr(cmdline, '\0', size);
    if(argv0_end)
        return argv0_end - cmdline;
    else
        return 0;
}

static const char* memchr_reverse(const char *p, int c, size_t size) {
    for(size_t i = 0; i < size; ++i) {
        if(p[size - 1 - i] == c)
            return p + i;
    }
    return NULL;
}

static const char* get_arg_by_index(const char *cmdline, size_t size, size_t index) {
    size_t current_index = 0;
    size_t cmdline_index = 0;
    while (index < size) {
        const char *arg_start = cmdline + cmdline_index;
        const char *arg_end = memchr(arg_start, '\0', size - cmdline_index);
        if(!arg_end)
            break;

        const size_t arg_len = arg_end - arg_start;
        if(current_index == index)
            return arg_start;

        cmdline_index += arg_len + 1;
        current_index += 1;
    }
    return NULL;
}

static bool memeql(const char *haystack, size_t haystack_size, const char *needle) {
    const size_t needle_size = strlen(needle);
    return haystack_size == needle_size && memcmp(haystack, needle, needle_size) == 0;
}

static bool starts_with(const char *haystack, size_t haystack_size, const char *needle) {
    const size_t needle_size = strlen(needle);
    return haystack_size >= needle_size && memcmp(haystack, needle, needle_size) == 0;
}

static int is_wine_binary(const char *argv0, size_t size) {
    const char *base = memchr_reverse(argv0, '/', size);
    base = base ? base + 1 : argv0;
    const size_t base_len = argv0 + size - base;
    return memeql(base, base_len, "wine") ||
           memeql(base, base_len, "wine64") ||
           memeql(base, base_len, "wine-preloader") ||
           memeql(base, base_len, "wine64-preloader");
}

static int has_game_arch_suffix(const char *argv0, size_t size) {
    static const char *suffixes[] = { ".x86_64", ".x64", ".x86" };
    for (int i = 0; i < 3; i++) {
        const size_t slen = strlen(suffixes[i]);
        if (size >= slen && memcmp(argv0 + size - slen, suffixes[i], slen) == 0)
            return 1;
    }
    return 0;
}

static void check_process(pid_t pid) {
    if (find_slot_by_pid(pid) >= 0) return;

    char    path[64];
    ssize_t env_n, cmd_n;
    bool is_steam_fossilize = false;
    bool is_stream_script = false;
    /* Proton launched outside steam has this while it doesn't have SteamAppId nor is the process called wine */
    bool has_wine_prefix_env = false;

    snprintf(path, sizeof(path), "/proc/%d/environ", pid);
    env_n = read_file(path, environ_buf, sizeof(environ_buf));

    if (env_n > 0) {
        is_steam_fossilize = env_get(environ_buf, env_n, "STEAM_FOSSILIZE_DUMP_PATH_READ_ONLY");
        is_stream_script = env_get(environ_buf, env_n, "STEAMSCRIPT_VERSION");
        has_wine_prefix_env = env_get(environ_buf, env_n, "WINEPREFIX");
        const char *appid = env_get(environ_buf, env_n, "SteamAppId");
        if (!appid) appid = env_get(environ_buf, env_n, "SteamGameId");
        if (appid && appid[0] >= '1' && appid[0] <= '9') {
            add_game(pid);
            return;
        }
    }

    snprintf(path, sizeof(path), "/proc/%d/cmdline", pid);
    cmd_n = read_file(path, cmdline_buf, sizeof(cmdline_buf));
    if(cmd_n <= 0)
        return;

    const char *argv0 = cmdline_buf;
    const size_t argv0_len = get_argv_len(cmdline_buf, cmd_n);

    const char *argv1 = get_arg_by_index(cmdline_buf, cmd_n, 1);
    const size_t argv1_len = argv1 ? get_argv_len(argv1, cmdline_buf + cmd_n - argv1) : 0;

    if (cmd_n > 0 && (is_wine_binary(argv0, argv0_len) || has_game_arch_suffix(argv0, argv0_len) || has_wine_prefix_env)
        && !is_steam_fossilize
        && !is_stream_script
        && !memeql(argv1, argv1_len, "--version")
        && !starts_with(argv0, argv0_len, "C:\\windows\\system32"))
    {
        // fprintf(stderr, "env: ");
        // write(STDOUT_FILENO, environ_buf, env_n);
        // fprintf(stderr, "\ncmdline: ");
        // write(STDOUT_FILENO, cmdline_buf, cmd_n);
        // fprintf(stderr, "\n");
        add_game(pid);
    }
}

static void handle_proc_event(const struct proc_event *ev) {
    switch (ev->what) {
    case PROC_EVENT_EXEC:
        check_process(ev->event_data.exec.process_tgid);
        break;
    case PROC_EVENT_EXIT:
        /* Only act when the whole process (not just a thread) exits */
        if (ev->event_data.exit.process_pid == ev->event_data.exit.process_tgid)
            remove_game(ev->event_data.exit.process_tgid);
        break;
    default:
        break;
    }
}

static void process_netlink_msg(const char *buf, size_t len) {
    const struct nlmsghdr *nl = (const struct nlmsghdr *)buf;
    for (; NLMSG_OK(nl, (unsigned int)len); nl = NLMSG_NEXT(nl, len)) {
        if (nl->nlmsg_type == NLMSG_NOOP  ||
            nl->nlmsg_type == NLMSG_ERROR ||
            nl->nlmsg_type == NLMSG_OVERRUN)
            continue;
        if (nl->nlmsg_len < NLMSG_HDRLEN + sizeof(struct cn_msg))
            continue;
        const struct cn_msg *cn = (const struct cn_msg *)NLMSG_DATA(nl);
        if (cn->id.idx != CN_IDX_PROC || cn->id.val != CN_VAL_PROC)
            continue;
        if (cn->len < sizeof(struct proc_event))
            continue;
        handle_proc_event((const struct proc_event *)cn->data);
    }
}

static int send_mcast_op(int sock, enum proc_cn_mcast_op op) {
    /* cn_msg ends with data[0], so we can't place fields after it in a struct.
     * Use a flat buffer and fill via pointers instead. */
    char buf[NLMSG_SPACE(sizeof(struct cn_msg) + sizeof(enum proc_cn_mcast_op))];
    memset(buf, 0, sizeof(buf));

    struct nlmsghdr *nl = (struct nlmsghdr *)buf;
    nl->nlmsg_len  = sizeof(buf);
    nl->nlmsg_type = NLMSG_DONE;
    nl->nlmsg_pid  = (unsigned int)getpid();

    struct cn_msg *cn = (struct cn_msg *)NLMSG_DATA(nl);
    cn->id.idx = CN_IDX_PROC;
    cn->id.val = CN_VAL_PROC;
    cn->len    = sizeof(enum proc_cn_mcast_op);
    memcpy(cn->data, &op, sizeof(op));

    return send(sock, buf, sizeof(buf), 0) < 0 ? -1 : 0;
}

static int setup_netlink(void) {
    int sock = socket(AF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR);
    if (sock < 0) {
        perror("socket(AF_NETLINK, SOCK_DGRAM, NETLINK_CONNECTOR)");
        return -1;
    }

    struct sockaddr_nl sa;
    memset(&sa, 0, sizeof(sa));
    sa.nl_family = AF_NETLINK;
    sa.nl_groups = CN_IDX_PROC;
    sa.nl_pid    = (unsigned int)getpid();

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

    if (send_mcast_op(sock, PROC_CN_MCAST_LISTEN) < 0) {
        perror("send PROC_CN_MCAST_LISTEN");
        close(sock);
        return -1;
    }

    return sock;
}

static void scan_existing_processes(void) {
    DIR *d = opendir("/proc");
    if (!d) return;

    struct dirent *ent;
    while ((ent = readdir(d)) != NULL) {
        const char *s = ent->d_name;
        if (*s < '1' || *s > '9') continue;
        pid_t pid = 0;
        for (; *s; s++) {
            if (*s < '0' || *s > '9') { pid = 0; break; }
            pid = pid * 10 + (*s - '0');
        }
        if (pid > 0) check_process(pid);
    }
    closedir(d);
}

int main(void) {
    memset(games, 0, sizeof(games));

    int sock = setup_netlink();
    if (sock < 0) return 1;

    scan_existing_processes();

    for (;;) {
        ssize_t n = recv(sock, recv_buf, sizeof(recv_buf), 0);
        if (n < 0) {
            if (errno == EINTR) continue;
            if (errno == ENOBUFS) { scan_existing_processes(); continue; }
            perror("recv");
            break;
        }
        process_netlink_msg(recv_buf, (size_t)n);
    }

    send_mcast_op(sock, PROC_CN_MCAST_IGNORE);
    close(sock);
    return 0;
}