/*
 * This file and its contents are supplied under the terms of the
 * Common Development and Distribution License ("CDDL"), version 1.0.
 * You may only use this file in accordance with the terms of version
 * 1.0 of the CDDL.
 *
 * A full copy of the text of the CDDL should have accompanied this
 * source.  A copy of the CDDL is also available via the Internet at
 * http://www.illumos.org/license/CDDL.
 *
 * Copyright 2023 Oxide Computer Company
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <err.h>
#include <assert.h>
#include <signal.h>
#include <sys/types.h>

#include <sys/vmm_dev.h>

const char *prog_name;

static void
usage(int exitcode)
{
        assert(prog_name != NULL);
        fprintf(stderr,
            "Usage: %s [-a add] [-r remove] [-q]\n"
            "\t-a <SZ> add SZ MiB to the reservoir\n"
            "\t-r <SZ> remove SZ MiB from the reservoir\n"
            "\t-s <SZ> set reservoir to SZ MiB, if possible\n"
            "\t-c <SZ> use SZ MiB chunks when performing resize ops\n"
            "\t-q query reservoir state\n", prog_name);
        exit(exitcode);
}

/*
 * Parse an input size of MiB to bytes.
 */
static bool
parse_size(const char *arg, size_t *resp)
{
        size_t res;

        errno = 0;
        res = strtoul(arg, NULL, 0);
        if (errno != 0) {
                return (false);
        }

        *resp = (res * 1024 * 1024);
        return (true);
}

static size_t
query_size(int fd)
{
        struct vmm_resv_query data;

        int res = ioctl(fd, VMM_RESV_QUERY, &data);
        if (res != 0) {
                err(EXIT_FAILURE, "Could not query reservoir sizing");
        }

        return (data.vrq_free_sz + data.vrq_alloc_sz);
}

static void
do_add(int fd, size_t sz, size_t chunk)
{
        const size_t cur = query_size(fd);
        struct vmm_resv_target target = {
                .vrt_target_sz = cur + sz,
                .vrt_chunk_sz = MIN(chunk, sz),
        };

        if (ioctl(fd, VMM_RESV_SET_TARGET, &target) != 0) {
                err(EXIT_FAILURE, "Could not add %zu bytes to reservoir", sz);
        }
}

static void
do_remove(int fd, size_t sz, size_t chunk)
{
        const size_t cur = query_size(fd);
        if (cur == 0) {
                /* Reservoir is already empty */
                return;
        }

        const size_t clamped_sz = MIN(sz, cur);
        struct vmm_resv_target target = {
                .vrt_target_sz = cur - clamped_sz,
                .vrt_chunk_sz = MIN(chunk, sz),
        };

        if (ioctl(fd, VMM_RESV_SET_TARGET, &target) != 0) {
                err(EXIT_FAILURE, "Could not remove %zu bytes from reservoir",
                    clamped_sz);
        }
}


bool caught_siginfo = false;

static void
siginfo_handler(int sig, siginfo_t *sip, void *ucp)
{
        caught_siginfo = true;
}

static void
do_set_target(int fd, size_t sz, size_t chunk)
{
        struct vmm_resv_target target = {
                .vrt_target_sz = sz,
                .vrt_chunk_sz = chunk,
        };

        struct sigaction sa = {
                .sa_sigaction = siginfo_handler,
                .sa_flags = SA_SIGINFO,
        };
        if (sigaction(SIGINFO, &sa, NULL) != 0) {
                err(EXIT_FAILURE, "Could not configure SIGINFO handler");
        }

        do {
                if (ioctl(fd, VMM_RESV_SET_TARGET, &target) != 0) {
                        if (errno != EINTR) {
                                err(EXIT_FAILURE,
                                    "Could not set reservoir size to %zu bytes",
                                    sz);
                        }

                        if (caught_siginfo) {
                                caught_siginfo = false;
                                (void) printf("Reservoir size: %zu MiB\n",
                                    target.vrt_result_sz / (1024 * 1024));
                        }
                }
        } while (target.vrt_result_sz != sz);
}

static void
do_query(int fd)
{
        struct vmm_resv_query data;
        int res;

        res = ioctl(fd, VMM_RESV_QUERY, &data);
        if (res != 0) {
                perror("Could not query reservoir info");
                return;
        }

        printf("Free MiB:\t%zu\n"
            "Allocated MiB:\t%zu\n"
            "Transient Allocated MiB:\t%zu\n"
            "Size limit MiB:\t%zu\n",
            data.vrq_free_sz / (1024 * 1024),
            data.vrq_alloc_sz / (1024 * 1024),
            data.vrq_alloc_transient_sz / (1024 * 1024),
            data.vrq_limit / (1024 * 1024));
}

int
main(int argc, char *argv[])
{
        int c;
        const char *opt_a = NULL, *opt_r = NULL, *opt_s = NULL;
        bool opt_q = false;
        int fd;

        prog_name = argv[0];

        uint_t resize_opts = 0;
        size_t chunk_sz = 0;
        while ((c = getopt(argc, argv, "a:r:s:c:qh")) != -1) {
                switch (c) {
                case 'a':
                        if (opt_a == NULL) {
                                resize_opts++;
                                opt_a = optarg;
                        }
                        break;
                case 'r':
                        if (opt_r == NULL) {
                                resize_opts++;
                                opt_r = optarg;
                        }
                        break;
                case 's':
                        if (opt_s == NULL) {
                                resize_opts++;
                                opt_s = optarg;
                        }
                        break;
                case 'c':
                        if (!parse_size(optarg, &chunk_sz)) {
                                warn("Invalid chunk size %s", optarg);
                                usage(EXIT_FAILURE);
                        }
                        break;
                case 'q':
                        opt_q = true;
                        break;
                case 'h':
                        usage(EXIT_SUCCESS);
                        break;
                default:
                        usage(EXIT_FAILURE);
                        break;
                }
        }

        if (optind < argc ||
            (resize_opts == 0 && !opt_q) || (resize_opts > 1)) {
                usage(EXIT_FAILURE);
        }

        fd = open(VMM_CTL_DEV, O_EXCL | O_RDWR);
        if (fd < 0) {
                perror("Could not open vmmctl");
                usage(EXIT_FAILURE);
        }

        if (opt_a != NULL) {
                size_t sz;

                if (!parse_size(opt_a, &sz)) {
                        warn("Invalid size %s", opt_a);
                        usage(EXIT_FAILURE);
                }

                do_add(fd, sz, chunk_sz);
        } else if (opt_r != NULL) {
                size_t sz;

                if (!parse_size(opt_r, &sz)) {
                        warn("Invalid size %s", opt_r);
                        usage(EXIT_FAILURE);
                }
                do_remove(fd, sz, chunk_sz);
        } else if (opt_s != NULL) {
                size_t sz;

                if (!parse_size(opt_s, &sz)) {
                        warn("Invalid size %s", opt_s);
                        usage(EXIT_FAILURE);
                }
                do_set_target(fd, sz, chunk_sz);
        } else if (opt_q) {
                do_query(fd);
        }

        (void) close(fd);
        return (0);
}