#define _XOPEN_SOURCE 500
#define _GNU_SOURCE
#include <stdio.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdbool.h>
#include <sys/types.h>
#include <signal.h>
#include <sys/time.h>
#include <sys/syscall.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <stdarg.h>

#define USEC_PER_SEC 1000000ULL

static long get_cguest_time(const char *buffer)
{
    const char *ptr;
    int space_count;

    for (ptr = buffer; *ptr && space_count != 42; ptr++) {
        if (*ptr == ' ') {
            space_count++;
        }
    }

    return strtol(ptr, NULL, 10);
}

static void tv_add(struct timeval *tv, suseconds_t usec)
{
    tv->tv_usec += usec;

    while (tv->tv_usec > USEC_PER_SEC) {
        tv->tv_sec += 1;
        tv->tv_usec -= USEC_PER_SEC;
    }
}

static int tv_cmp(struct timeval *lhs, struct timeval *rhs)
{
    if (lhs->tv_sec == rhs->tv_sec) {
        if (lhs->tv_usec < rhs->tv_usec) {
            return -1;
        } else if (lhs->tv_usec > rhs->tv_usec) {
            return 1;
        }
        return 0;
    } else if (lhs->tv_sec < rhs->tv_sec) {
        return -1;
    } else if (lhs->tv_sec > rhs->tv_sec) {
        return 1;
    }
    return 0;
}

static void write_all(int fd, const void *buffer, size_t buffer_len)
{
    size_t offset = 0;

    while (offset < buffer_len) {
        ssize_t len;

        len = write(fd, buffer + offset, buffer_len - offset);
        if (len > 0) {
            offset += len;
        }
    }
}

static void read_reply(int fd, char *buffer, size_t buffer_len)
{
    size_t offset = 0;

    while (offset < buffer_len) {
        ssize_t len;

        len = read(fd, buffer + offset, buffer_len - offset);
        if (len > 0) {
            offset += len;
        }
        if (offset > 8 &&
            memcmp("\n(qemu) ", buffer + (offset - 8), 8) == 0) {
            char *ptr;
            buffer[offset - 8] = 0;
            ptr = strchr(buffer, '\n');
            if (ptr == NULL) {
                buffer[0] = 0;
            } else {
                memmove(buffer, ptr + 1, offset - (ptr - buffer) - 1);
            }
            return;
        }
    }
}

static int monitor_fd;

static void monitor_command(const char *fmt, ...)
{
    char buffer[256];
    va_list ap;
    size_t len;
    
    va_start(ap, fmt);
    len = vsnprintf(buffer, sizeof(buffer), fmt, ap);
    va_end(ap);

    write_all(monitor_fd, buffer, len);
    write_all(monitor_fd, "\n", 1);
    read_reply(monitor_fd, buffer, sizeof(buffer));
}

static void monitor_command_response(char *rsp, size_t rsp_len,
                                     const char *fmt, ...)
{
    char buffer[256];
    va_list ap;
    size_t len;
    
    va_start(ap, fmt);
    len = vsnprintf(buffer, sizeof(buffer), fmt, ap);
    va_end(ap);

    write_all(monitor_fd, buffer, len);
    write_all(monitor_fd, "\n", 1);
    read_reply(monitor_fd, rsp, rsp_len);
}

static int vm_running = 1;

static void guest_start(int vcpu)
{
    if (!vm_running) {
        monitor_command("cpu_start %d", vcpu);
    }
    vm_running = 1;
}

static void guest_stop(int vcpu)
{
    if (vm_running) {
        monitor_command("cpu_stop %d", vcpu);
    }
    vm_running = 0;
}

static int find_pid(char *buffer, int vcpu)
{
    char *ptr = buffer;
    int i;

    for (i = 0; ptr && i < vcpu; i++) {
        ptr = strchr(ptr, '\n');
        if (ptr) {
            ptr++;
        }
    }

    if (ptr) {
        ptr = strstr(ptr, "thread_id=");
        if (ptr) {
            ptr += 10;
            return atoi(ptr);
        }
    }

    return 0;
}

int main(int argc, char **argv)
{
    int fd, pid, vcpu;
    char buffer[1024];
    long ticks_per_sec;
    long cguest_time_last = 0;
    struct timeval period_end;
    long cguest_ticks;
    long entitlement;
    long period;
    struct sockaddr_un addr;

    if (argc != 4) {
        fprintf(stderr, "Missing arguments\n");
        return 1;
    }

    vcpu = atoi(argv[1]);
    /* FIXME hack, does guest time get scaled with vcpu count? */
    entitlement = atoi(argv[2]) * 2;
    period = atoi(argv[3]);

    monitor_fd = socket(PF_UNIX, SOCK_STREAM, 0);
    addr.sun_family = AF_UNIX;
    snprintf(addr.sun_path, 108, "/tmp/monitor.sock");

    if (connect(monitor_fd, (struct sockaddr *)&addr, sizeof(addr)) == -1) {
        return 1;
    }

    read_reply(monitor_fd, buffer, sizeof(buffer));
    monitor_command_response(buffer, sizeof(buffer), "info cpus");
    pid = find_pid(buffer, vcpu);

    ticks_per_sec = sysconf(_SC_CLK_TCK);
    entitlement = (entitlement * ticks_per_sec) / 1000;
    period *= 1000;

    snprintf(buffer, sizeof(buffer), "/proc/%d/stat", pid);
    fd = open(buffer, O_RDONLY);

    gettimeofday(&period_end, NULL);
    tv_add(&period_end, period);
    cguest_ticks = 0;

    while (1) {
        long cguest_time_now;
        struct timeval tv_now;
        ssize_t len;

        gettimeofday(&tv_now, NULL);
        len = pread(fd, buffer, sizeof(buffer) - 1, 0);
        buffer[len] = 0;
        cguest_time_now = get_cguest_time(buffer);

        if (cguest_time_last) {
            cguest_ticks += cguest_time_now - cguest_time_last;

            if (tv_cmp(&tv_now, &period_end) < 0) {
                if (cguest_ticks >= entitlement) {
                    guest_stop(vcpu);
                    cguest_ticks = 0;
                }
            } else {
                guest_start(vcpu);
                cguest_ticks = 0;
                tv_add(&tv_now, period);
                period_end = tv_now;
            }
        }

        cguest_time_last = cguest_time_now;
        usleep(10000); // 10ms
    }

    close(fd);

    return 0;
}