#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <limits.h>
#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdint.h>
#include <sys/stat.h>
#include <sys/capability.h>

#define KMS_SERVER_PROXY_FILEPATH "/var/lib/flatpak/app/com.dec05eba.gpu_screen_recorder/current/active/files/bin/kms-server-proxy"
#define GSR_KMS_SERVER_FILEPATH "/var/lib/flatpak/app/com.dec05eba.gpu_screen_recorder/current/active/files/bin/gsr-kms-server"
#define GSR_GLOBAL_HOTKEYS_FILEPATH "/var/lib/flatpak/app/com.dec05eba.gpu_screen_recorder/current/active/files/bin/gsr-global-hotkeys"

static bool readlink_realpath(const char *filepath, char *buffer) {
    char symlinked_path[PATH_MAX];
    ssize_t bytes_written = readlink(filepath, symlinked_path, sizeof(symlinked_path) - 1);
    if(bytes_written == -1 && errno == EINVAL) {
        /* Not a symlink */
        strncpy(symlinked_path, filepath, sizeof(symlinked_path));
    } else if(bytes_written == -1) {
        return false;
    } else {
        symlinked_path[bytes_written] = '\0';
    }

    if(!realpath(symlinked_path, buffer))
        return false;

    return true;
}

static bool file_has_capabilities(const char *filepath, const cap_value_t *caps, int num_caps) {
    cap_t cap = cap_get_file(filepath);
    if(!cap)
        return false;

    bool has_caps = true;
    for(int i = 0; i < num_caps; ++i) {
        cap_flag_value_t res = CAP_CLEAR;
        cap_get_flag(cap, caps[i], CAP_PERMITTED, &res);
        if(res != CAP_SET) {
            has_caps = false;
            break;
        }
    }

    cap_free(cap);
    return has_caps;
}

static bool file_set_capabilities(const char *filepath, const cap_value_t *caps_to_set, int num_caps_to_set) {
    cap_t cap = cap_get_file(filepath);
    if(!cap && errno != ENODATA)
        return false;

    if(!cap) {
        cap = cap_init();
        if(!cap)
            return false;
    }

    int res = 0;
    res |= cap_set_flag(cap, CAP_EFFECTIVE, num_caps_to_set, caps_to_set, CAP_SET);
    res |= cap_set_flag(cap, CAP_PERMITTED, num_caps_to_set, caps_to_set, CAP_SET);
    /*cap_set_flag(cap, CAP_INHERITABLE, num_caps_to_set, caps_to_set, CAP_SET);*/
    res |= cap_set_file(filepath, cap);
    cap_free(cap);
    return res == 0;
}

static const char* get_column(const char *str, int size, int *column_size, int column) {
    *column_size = 0;
    int current_column = 0;
    int i = 0;
    for(; column > 0 && i < size; ++i) {
        if(str[i] == ' ') {
            ++current_column;
            if(current_column == column) {
                ++i;
                break;
            }
        }
    }

    const int column_start = i;
    for(; i < size; ++i) {
        if(str[i] == ' ') {
            *column_size = i - column_start;
            return str + column_start;
        }
    }

    *column_size = size;
    return str;
}

static bool mount_has_flag(const char *flags, int flags_size, const char *flag_to_find) {
    if(flags_size == 0)
        return false;

    const int flag_to_find_size = strlen(flag_to_find);
    int offset = 0;
    for(;;) {
        const char *p = memmem(flags + offset, flags_size - offset, flag_to_find, flag_to_find_size);
        if(!p)
            return false;

        offset = p - flags;
        const char prev_char = offset > 0 ? flags[offset - 1] : ',';
        const char next_char = offset + flag_to_find_size < flags_size ? flags[offset + flag_to_find_size] : ',';
        if((prev_char == ',' || prev_char == '\n') && (next_char == ',' || next_char == '\n'))
            return true;
    }

    return false;
}

static bool mount_has_setuid_permission(const char *path) {
    FILE *f = fopen("/proc/mounts", "rb");
    if(!f) {
        fprintf(stderr, "Error: failed to open /proc/mounts\n");
        // Assume it's ok and just continue
        return true;
    }

    // Assume there is permission if it can't be found
    bool has_permission = true;
    const int path_size = strlen(path);
    char line[1024];
    while(fgets(line, sizeof(line), f)) {
        const int line_size = strlen(line);
        int mount_point_size = 0;
        const char *mount_point = get_column(line, line_size, &mount_point_size, 1);
        if(!mount_point)
            continue;

        if(mount_point && mount_point_size == path_size && memcmp(mount_point, path, mount_point_size) == 0) {
            int flags_size = 0;
            const char *flags = get_column(line, line_size, &flags_size, 3);
            if(!flags)
                continue;

            const bool has_rw = mount_has_flag(flags, flags_size, "rw");
            const bool has_nosuid = mount_has_flag(flags, flags_size, "nosuid");
            has_permission = has_rw && !has_nosuid;
            break;
        }
    }

    fclose(f);
    return has_permission;
}

static bool create_local_kms_server_proxy_directory(const char *home) {
    char path[PATH_MAX];
    int err;

    const char *paths[] = { ".local", ".local/share", ".local/share/gpu-screen-recorder", NULL };
    for(size_t i = 0; paths[i]; ++i) {
        const char *path_part = paths[i];
        snprintf(path, sizeof(path), "%s/%s", home, path_part);
        err = mkdir(path, S_IRWXU);
        if(err == -1 && errno != EEXIST)
            return false;
    }

    return true;
}

static bool copy_file_atomic(const char *source_path, const char *dest_path) {
    int in_fd = -1;
    int out_fd = -1;
    bool res = false;
    char tmp_filepath[PATH_MAX];
    snprintf(tmp_filepath, sizeof(tmp_filepath), "%s.tmp", dest_path);

    in_fd = open(source_path, O_RDONLY);
    if(in_fd == -1) {
        fprintf(stderr, "Error: copy_file_atomic: failed to open %s, error: %s\n", source_path, strerror(errno));
        goto done;
    }

    out_fd = open(tmp_filepath, O_WRONLY | O_CREAT, 0755);
    if(out_fd == -1) {
        fprintf(stderr, "Error: copy_file_atomic: failed to create %s, error: %s\n", tmp_filepath, strerror(errno));
        goto done;
    }

    uint8_t buffer[8192];
    for(;;) {
        const ssize_t bytes_read = read(in_fd, buffer, sizeof(buffer));
        if(bytes_read == -1) {
            fprintf(stderr, "Error: copy_file_atomic: failed to read data from %s, error: %s\n", source_path, strerror(errno));
            goto done;
        } else if(bytes_read == 0) {
            break;
        }

        ssize_t bytes_written_total = 0;
        while(bytes_written_total < bytes_read) {
            const ssize_t bytes_written = write(out_fd, buffer + bytes_written_total, bytes_read - bytes_written_total);
            if(bytes_written == -1) {
                fprintf(stderr, "Error: copy_file_atomic: failed to write data to %s, error: %s\n", tmp_filepath, strerror(errno));
                goto done;
            }
            bytes_written_total += bytes_written;
        }
    }

    res = true;

    done:
    if(in_fd)
        close(in_fd);
    if(out_fd)
        close(out_fd);

    if(res) {
        res = rename(tmp_filepath, dest_path) == 0;
        if(!res)
            fprintf(stderr, "Error: copy_file_atomic: Failed to rename from %s to %s, error: %s\n", strerror(errno), tmp_filepath, dest_path);
    } else {
        remove(tmp_filepath);
    }

    return res;
}

static bool gsr_files_set_permissions_and_capabilities(const char *kms_server_proxy_local_filepath) {
    /* owner: read/write/execute, group: read/execute, public: read/execute */
    if(chmod(kms_server_proxy_local_filepath, 0755) != 0) {
        fprintf(stderr, "Error: failed to set %s permissions\n", kms_server_proxy_local_filepath);
        return false;
    }

    if(chown(kms_server_proxy_local_filepath, 0, 0) != 0) {
        fprintf(stderr, "Error: failed to set %s ownership\n", kms_server_proxy_local_filepath);
        return false;
    }

    if(!file_set_capabilities(kms_server_proxy_local_filepath, (const cap_value_t[]){ CAP_SYS_ADMIN, CAP_SETUID }, 2)) {
        fprintf(stderr, "Error: failed to set kms-server-proxy capabilities\n");
        return false;
    }

    return true;
}

/* |gsr_global_hotkeys_local_filepath| can be NULL */
static bool setup_local_gsr_files(const char *user_homepath, const char *kms_server_proxy_home) {
    if(!create_local_kms_server_proxy_directory(user_homepath)) {
        fprintf(stderr, "Error: failed to create ~/.local/share/gpu-screen-recorder directory\n");
        return false;
    }

    remove(kms_server_proxy_home);
    if(!copy_file_atomic(KMS_SERVER_PROXY_FILEPATH, kms_server_proxy_home)) {
        fprintf(stderr, "Error: failed to copy kms-server-proxy to %s\n", kms_server_proxy_home);
        return false;
    }

    return true;
}

static bool get_local_kms_server_proxy_filepath(char *output, size_t output_size, const char *user_homepath) {
    if(mount_has_setuid_permission(user_homepath)) {
        snprintf(output, output_size, "%s/.local/share/gpu-screen-recorder/kms-server-proxy-2", user_homepath);
        return true;
    } else if(mount_has_setuid_permission("/")) {
        snprintf(output, output_size, "/usr/bin/kms-server-proxy-2");
        return true;
    } else {
        fprintf(stderr, "Error: neither %s nor / has both read-write permission and setuid permission\n", user_homepath);
        return false;
    }
}

static int setup_gsr_ui(const char *user_homepath) {
    if(!user_homepath) {
        user_homepath = getenv("HOME");
        if(!user_homepath)
            user_homepath = "/tmp";
    }

    char kms_server_proxy_local_filepath[PATH_MAX];
    if(!get_local_kms_server_proxy_filepath(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), user_homepath))
        return 1;

    /* Weird-ass distros like openSUSE only allow pkexec for files in $HOME!!! */
    char kms_server_proxy_home[PATH_MAX];
    snprintf(kms_server_proxy_home, sizeof(kms_server_proxy_home), "%s/kms-server-proxy", user_homepath);

    if(geteuid() == 0) { /* is current user root? */
        remove(kms_server_proxy_home);
        remove(kms_server_proxy_local_filepath);
        if(!copy_file_atomic(KMS_SERVER_PROXY_FILEPATH, kms_server_proxy_local_filepath)) {
            fprintf(stderr, "Error: failed to copy kms-server-proxy to %s\n", kms_server_proxy_local_filepath);
            return 1;
        }

        if(!gsr_files_set_permissions_and_capabilities(kms_server_proxy_local_filepath))
            return 1;

        return 0;
    } else {
        if(!setup_local_gsr_files(user_homepath, kms_server_proxy_home))
            return 1;

        const char *args[] = { "pkexec", kms_server_proxy_home, "setup-gsr-ui", user_homepath, NULL };
        const int result = execvp(args[0], (char *const*)args);
        perror("pkexec");
        remove(kms_server_proxy_home);
        return result;
    }
}

static bool is_setup(const char *user_homepath) {
    if(!user_homepath) {
        user_homepath = getenv("HOME");
        if(!user_homepath)
            user_homepath = "/tmp";
    }

    char kms_server_proxy_local_filepath[PATH_MAX];
    if(!get_local_kms_server_proxy_filepath(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), user_homepath))
        return false;

    if(!file_has_capabilities(kms_server_proxy_local_filepath, (const cap_value_t[]){ CAP_SYS_ADMIN, CAP_SETUID }, 2)) {
        fprintf(stderr, "Error: kms-server-proxy (%s) is either not installed or is missing capabilities\n", kms_server_proxy_local_filepath);
        return false;
    }

    return true;
}

static int launch_gsr_kms_server(const char *initial_socket_path, const char *card_path, const char *user_homepath) {
    char self_path[PATH_MAX];
    if(!readlink_realpath("/proc/self/exe", self_path)) {
        fprintf(stderr, "failed to resolve /proc/self/exe\n");
        return 1;
    }

    char kms_server_proxy_local_filepath[PATH_MAX];
    if(!get_local_kms_server_proxy_filepath(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), user_homepath))
        return 1;

    /* Weird-ass distros like openSUSE only allow pkexec for files in $HOME!!! */
    char kms_server_proxy_home[PATH_MAX];
    snprintf(kms_server_proxy_home, sizeof(kms_server_proxy_home), "%s/kms-server-proxy", user_homepath);

    if(file_has_capabilities(kms_server_proxy_local_filepath, (const cap_value_t[]){ CAP_SYS_ADMIN, CAP_SETUID }, 2)) {
        /* Need to resolve kms_server_proxy_local_filepath because /home can be a symlink to another location */
        char kms_server_proxy_local_filepath_full[PATH_MAX];
        if(!readlink_realpath(kms_server_proxy_local_filepath, kms_server_proxy_local_filepath_full)) {
            fprintf(stderr, "failed to resolve %s\n", kms_server_proxy_local_filepath);
            return 1;
        }

        /*
            Run cached ~/.local/share/gpu-screen-recorder/kms-server-proxy which has sys admin capability.
            The one in flatpak location gets capabilities overwritten on flatpak update.
        */
        if(strcmp(self_path, kms_server_proxy_local_filepath_full) != 0) {
            const char *args[] = { kms_server_proxy_local_filepath_full, initial_socket_path, card_path, user_homepath, NULL };
            return execv(args[0], (char *const*)args);
        }

        if(setuid(0) == -1) {
            fprintf(stderr, "Error: failed to switch to root user to launch gsr-kms-server\n");
            return 1;
        }

        const char *args[] = { GSR_KMS_SERVER_FILEPATH, initial_socket_path, card_path, NULL };
        return execvp(args[0], (char *const*)args);
    } else if(geteuid() == 0) { /* is current user root? */
        remove(kms_server_proxy_home);
        remove(kms_server_proxy_local_filepath);
        if(!copy_file_atomic(KMS_SERVER_PROXY_FILEPATH, kms_server_proxy_local_filepath)) {
            fprintf(stderr, "Error: failed to copy kms-server-proxy to %s\n", kms_server_proxy_local_filepath);
            return 1;
        }

        if(!gsr_files_set_permissions_and_capabilities(kms_server_proxy_local_filepath))
            return 1;

        const char *args[] = { GSR_KMS_SERVER_FILEPATH, initial_socket_path, card_path, NULL };
        return execv(args[0], (char *const*)args);
    } else {
        if(!setup_local_gsr_files(user_homepath, kms_server_proxy_home))
            return 1;

        const char *args[] = { "pkexec", kms_server_proxy_home, initial_socket_path, card_path, user_homepath, NULL };
        const int result = execvp(args[0], (char *const*)args);
        perror("pkexec");
        remove(kms_server_proxy_home);
        return result;
    }
}

static int launch_gsr_global_hotkeys(char **argv) {
    char self_path[PATH_MAX];
    if(!readlink_realpath("/proc/self/exe", self_path)) {
        fprintf(stderr, "failed to resolve /proc/self/exe\n");
        return 1;
    }

    const char *user_homepath = argv[2];

    char kms_server_proxy_local_filepath[PATH_MAX];
    if(!get_local_kms_server_proxy_filepath(kms_server_proxy_local_filepath, sizeof(kms_server_proxy_local_filepath), user_homepath))
        return 1;

    if(!file_has_capabilities(kms_server_proxy_local_filepath, (const cap_value_t[]){ CAP_SETUID }, 1)) {
        fprintf(stderr, "Error: kms-server-proxy (%s) is either not installed or is missing capabilities\n", kms_server_proxy_local_filepath);
        return 1;
    }

    /* Need to resolve kms_server_proxy_local_filepath because /home can be a symlink to another location */
    char kms_server_proxy_local_filepath_full[PATH_MAX];
    if(!readlink_realpath(kms_server_proxy_local_filepath, kms_server_proxy_local_filepath_full)) {
        fprintf(stderr, "failed to resolve %s\n", kms_server_proxy_local_filepath);
        return 1;
    }

    /*
        Run cached ~/.local/share/gpu-screen-recorder/kms-server-proxy which has sys admin capability.
        The one in flatpak location gets capabilities overwritten on flatpak update.
    */
    if(strcmp(self_path, kms_server_proxy_local_filepath_full) != 0) {
        argv[0] = kms_server_proxy_local_filepath_full;
        return execv(argv[0], argv);
    }

    if(setuid(0) == -1) {
        fprintf(stderr, "Error: failed to switch to root user to launch gsr-global-hotkeys\n");
        return 1;
    }

    argv[2] = GSR_GLOBAL_HOTKEYS_FILEPATH;
    const int result = execv(argv[2], argv + 2);
    perror(argv[2]);
    return result;
}

static void usage(void) {
    fprintf(stderr, "usage alt.1: kms-server-proxy <initial_socket_path> <card_path> <user_homepath>\n");
    fprintf(stderr, "usage alt.2: kms-server-proxy setup-gsr-ui [user_homepath]\n");
    fprintf(stderr, "usage alt.3: kms-server-proxy launch-gsr-global-hotkeys <user_homepath> [args...]\n");
    fprintf(stderr, "usage alt.4: kms-server-proxy is-setup [user_homepath]\n");
    exit(1);
}

int main(int argc, char **argv) {
    if(argc < 2)
        usage();

    if(strcmp(argv[1], "setup-gsr-ui") == 0) {
        const char *user_homepath = argc == 3 ? argv[2] : NULL;
        return setup_gsr_ui(user_homepath);
    } else if(strcmp(argv[1], "launch-gsr-global-hotkeys") == 0) {
        if(argc >= 3) {
            return launch_gsr_global_hotkeys(argv);
        } else {
            fprintf(stderr, "Error: missing user_homepath argument to launch-gsr-global-hotkeys\n");
            usage();
        }
    } else if(strcmp(argv[1], "is-setup") == 0) {
        const char *user_homepath = argc == 3 ? argv[2] : NULL;
        const bool s = is_setup(user_homepath);
        return s ? 0 : 1;
    } else if(argc == 4) {
        return launch_gsr_kms_server(argv[1], argv[2], argv[3]);
    } else {
        fprintf(stderr, "error: invalid option \"%s\".\n", argv[1]);
        usage();
    }
}