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

#include <sys/socket.h>
#include <sys/un.h>

struct msghdr mk_msghdr();
int ch_connect(const char*, const char*);
ssize_t send_fd(int, int);

#define _WRITE_CONFIRM(fd, buf, buflen) {if (write((fd), (buf), (buflen)) != (buflen)) { perror("ch-proxy/write/partial write"); exit(EXIT_FAILURE); }}

void print_usage() {
    fprintf(stderr, "%s",
            "Usage:\n"
            "\tch-proxy uvm/$USER_VM_NAME [PORT]\n"
            "\tch-proxy uuvm/$VM_NAME [PORT]\n"
            "\tch-proxy vsock-mux%$PATH [PORT]\n");
}

char *extract_vsock_mux(const char *host_string) {
    const char PREFIX[] = "vsock-mux%";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;
    if (strncmp(host_string, PREFIX, PREFIX_LEN) == 0) {
        return strdup(&host_string[PREFIX_LEN]);
    }
    return NULL;
}

char *extract_uvm(const char *host_string) {
    const char PREFIX[] = "uuvm/";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;

    const char *home = getenv("HOME") != NULL ? getenv("HOME") : "./.";

    if (strncmp(host_string, PREFIX, PREFIX_LEN) != 0) {
        return NULL;
    }

    char *result;
    if (asprintf(&result, "%s/uvms/%s/CONNECT.sock", home, &host_string[PREFIX_LEN]) == -1) {
        perror("ch-proxy/extract_uvm");
        exit(EXIT_FAILURE);
    }

    return result;
}

char *extract_muvm(const char *host_string) {
    const char PREFIX[] = "uvm/";
    const ssize_t PREFIX_LEN = sizeof(PREFIX) - 1;


    if (strncmp(host_string, PREFIX, PREFIX_LEN) != 0) {
        return NULL;
    }

    char *result;
    if (asprintf(&result, "/var/lib/microvms/%s/CONNECT.sock", &host_string[PREFIX_LEN]) == -1) {
        perror("ch-proxy/extract_muvm");
        exit(EXIT_FAILURE);
    }

    return result;
}

int main(int argc, char** argv) {
    if (!(2 <= argc && argc <= 3)) {
        fprintf(stderr, "%s: Wrong # of arguments: %d\n", argv[0], argc);
        print_usage();
        return 1;
    }

    const char PORT_DEFAULT[] = "22";

    const char *ssh_host = argv[1];
    const char *port_string = argc == 3 ? argv[2] : PORT_DEFAULT;

    char *path_un;
    if ((path_un = extract_uvm(ssh_host)) != NULL) {
    } else if ((path_un = extract_muvm(ssh_host)) != NULL) {
    } else if ((path_un = extract_vsock_mux(ssh_host)) != NULL) {
    } else {
        fprintf(stderr, "ch-proxy/main: unexpected host stirng format: %s\n", ssh_host);
        print_usage();
        return EXIT_FAILURE;
    }

    const int s = ch_connect(path_un, port_string);
    if (s == -1) {
        perror("ssh-vsock-proxy/main/ch_connect");
        return EXIT_FAILURE;
    };

    if (send_fd(1, s) == -1) {
        perror("ssh-vsock-proxy/main/send_fd");
        return EXIT_FAILURE;
    }
    return 0;
}

struct msghdr mk_msghdr() {
    struct msghdr msg;
    memset(&msg, 0, sizeof(msg));

    return msg;
}

int ch_connect(const char *path, const char *port) {
    int s = socket(AF_UNIX, SOCK_STREAM, 0);
    if (s == -1) {
        perror("ch-proxy/ch_connect/bind");
        return -1;
    }

    struct sockaddr_un a;
    a.sun_family = AF_UNIX;
    strcpy(a.sun_path, path);

    if (connect(s, (struct sockaddr*) &a, sizeof(a)) == -1) {
        char *err_msg;
        if (asprintf(&err_msg, "ch-proxy/ch_connect/connect(\"%s\")", path) == -1) {
            perror("ch-proxy/ch_connect/while printing connect() error");
            exit(EXIT_FAILURE);
        }
        perror(err_msg);
        exit(EXIT_FAILURE);
    }

    const char CMD[] = "CONNECT ";
    const int CMD_LEN = sizeof(CMD) - 1;

    _WRITE_CONFIRM(s, CMD, CMD_LEN);
    _WRITE_CONFIRM(s, port, strlen(port));
    _WRITE_CONFIRM(s, "\n", 1);

    return s;
}

ssize_t send_fd(int dst_fd, int fd) {
    struct msghdr msg = mk_msghdr();

    /* openssh expects to receive a dummy length=1 iovec? */
    char ch;
    struct iovec vec;
    vec.iov_base = &ch;
    vec.iov_len = 1;
    msg.msg_iov = &vec;
    msg.msg_iovlen = 1;

    union {
        struct cmsghdr align;
        char buf[CMSG_SPACE(sizeof(int))];
    } u;

    msg.msg_control = u.buf;
    msg.msg_controllen = sizeof(u.buf);

    struct cmsghdr *cmptr;
    cmptr = CMSG_FIRSTHDR(&msg);

    if (cmptr == NULL) {
        fprintf(stderr, "ch-proxy/send_fd/CMSG_FIRSTHDR: failed to initialize msg_control\n");
        exit(EXIT_FAILURE);
    }

    cmptr->cmsg_len = CMSG_LEN(sizeof(int));
    cmptr->cmsg_level = SOL_SOCKET;
    cmptr->cmsg_type = SCM_RIGHTS;
    *((int*) CMSG_DATA(cmptr)) = fd;

    return (sendmsg(dst_fd, &msg, 0));
}