// mkvirtualdrive.cpp

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

#include <Path.h>

#include "virtualdrive.h"

extern "C" const char * const*__libc_argv;

const char *kUsage =
"Usage: %s [ --install ] [ --size <size> ] <file>\n"
"       %s --uninstall  <device>\n"
"       %s --halt <device> (dangerous!)\n"
"       %s --info  <device>\n"
"       %s ( --help | -h )\n"
;

// print_usage
static
void
print_usage(bool error = false)
{
        const char *name = "mkvirtualdrive";
        BPath path;
        if (path.SetTo(__libc_argv[0]) == B_OK)
                name = path.Leaf();
        fprintf((error ? stderr : stdout), kUsage, name, name, name, name);
}


// parse_size
static
bool
parse_size(const char *str, off_t *_size)
{
        if (!str || !_size)
                return false;
        int32 len = strlen(str);
        char buffer[22];
        if (len == 0 || len + 1 > (int32)sizeof(buffer))
                return false;
        strcpy(buffer, str);
        // check the last char for one of the suffixes
        off_t suffix = 1;
        if (!isdigit(buffer[len - 1])) {
                switch (buffer[len - 1]) {
                        case 'K':
                                suffix = 1LL << 10;
                                break;
                        case 'M':
                                suffix = 1LL << 20;
                                break;
                        case 'G':
                                suffix = 1LL << 30;
                                break;
                        case 'T':
                                suffix = 1LL << 40;
                                break;
                        case 'P':
                                suffix = 1LL << 50;
                                break;
                        case 'E':
                                suffix = 1LL << 60;
                                break;
                        default:
                                return false;
                }
                buffer[len - 1] = '\0';
                len--;
                if (len == 0)
                        return false;
        }
        // all other chars must be digits
        for (int i = 0; i < len; i++) {
                if (!isdigit(buffer[i]))
                        return false;
        }
        off_t size = atoll(buffer);
        *_size = size * suffix;
        // check for overflow
        if (*_size / suffix != size)
                return false;
        return true;
}


// install_file
status_t
install_file(const char *file, off_t size)
{
        // open the control device
        int fd = open(VIRTUAL_DRIVE_CONTROL_DEVICE, O_RDONLY);
        if (fd < 0) {
                fprintf(stderr, "Failed to open control device: %s\n",
                                strerror(errno));
                return errno;
        }

        // set up the info
        virtual_drive_info info;
        info.magic = VIRTUAL_DRIVE_MAGIC;
        info.drive_info_size = sizeof(info);

        strcpy(info.file_name, file);
        if (size >= 0) {
                info.use_geometry = true;
                // fill in the geometry
                // default to 512 bytes block size
                uint32 blockSize = 512;
                // Optimally we have only 1 block per sector and only one head.
                // Since we have only a uint32 for the cylinder count, this won't work
                // for files > 2TB. So, we set the head count to the minimally possible
                // value.
                off_t blocks = size / blockSize;
                uint32 heads = (blocks + ULONG_MAX - 1) / ULONG_MAX;
                if (heads == 0)
                        heads = 1;
                info.geometry.bytes_per_sector = blockSize;
            info.geometry.sectors_per_track = 1;
            info.geometry.cylinder_count = blocks / heads;
            info.geometry.head_count = heads;
            info.geometry.device_type = B_DISK; // TODO: Add a new constant.
            info.geometry.removable = false;
            info.geometry.read_only = false;
            info.geometry.write_once = false;
        } else {
                info.use_geometry = false;
        }

        // issue the ioctl
        status_t error = B_OK;
        if (ioctl(fd, VIRTUAL_DRIVE_REGISTER_FILE, &info) != 0) {
                error = errno;
                fprintf(stderr, "Failed to install device: %s\n", strerror(error));
        } else {
                printf("File `%s' registered as device `%s'.\n", file, info.device_name);
        }
        // close the control device
        close(fd);
        return error;
}


// uninstall_file
status_t
uninstall_file(const char *device, bool immediately)
{
        // open the device
        int fd = open(device, O_RDONLY);
        if (fd < 0) {
                fprintf(stderr, "Failed to open device `%s': %s\n",
                                device, strerror(errno));
                return errno;
        }

        // issue the ioctl
        status_t error = B_OK;
        if (ioctl(fd, VIRTUAL_DRIVE_UNREGISTER_FILE, immediately) != 0) {
                error = errno;
                fprintf(stderr, "Failed to uninstall device: %s\n", strerror(error));
        }
        // close the control device
        close(fd);
        return error;
}


status_t
info_for_device(const char *device)
{
        // open the device
        int fd = open(device, O_RDONLY);
        if (fd < 0) {
                fprintf(stderr, "Failed to open device `%s': %s\n",
                                device, strerror(errno));
                return errno;
        }

        // set up the info
        virtual_drive_info info;
        info.magic = VIRTUAL_DRIVE_MAGIC;
        info.drive_info_size = sizeof(info);

        // issue the ioctl
        status_t error = B_OK;
        if (ioctl(fd, VIRTUAL_DRIVE_GET_INFO, &info) != 0) {
                error = errno;
                fprintf(stderr, "Failed to get device info: %s\n", strerror(error));
        } else {
                printf("Device `%s' points to file `%s'.\n", info.device_name, info.file_name);
                off_t size = (off_t)info.geometry.bytes_per_sector
                                * info.geometry.sectors_per_track
                                * info.geometry.cylinder_count
                                * info.geometry.head_count;
                printf("\tdisk size is %lld bytes (%g MB)\n", size, size / (1024.0 * 1024));
                if (info.halted)
                        printf("\tdevice is currently halted.\n");
        }
        // close the control device
        close(fd);
        return error;
}


// main
int
main(int argc, const char **argv)
{
        status_t error = B_OK;
        int argIndex = 1;
        enum { INSTALL, UNINSTALL, INFO } mode = INSTALL;
        bool immediately = false;
        off_t size = -1;

        // parse options
        for (; error == B_OK && argIndex < argc
                   && argv[argIndex][0] == '-'; argIndex++) {
                const char *arg = argv[argIndex];
                if (arg[1] == '-') {
                        // "--" option
                        arg += 2;
                        if (!strcmp(arg, "install")) {
                                mode = INSTALL;
                        } else if (!strcmp(arg, "size")) {
                                argIndex++;
                                if (argIndex >= argc) {
                                        fprintf(stderr, "Parameter expected for `--%s'.\n", arg);
                                        print_usage(1);
                                        return 1;
                                }
                                if (!parse_size(argv[argIndex], &size)) {
                                        fprintf(stderr, "Parameter for `--%s' must be of the form "
                                                        "`<number>[K|M|G|T|P|E]'.\n", arg);
                                        print_usage(1);
                                        return 1;
                                }
                        } else if (!strcmp(arg, "uninstall")) {
                                mode = UNINSTALL;
                        } else if (!strcmp(arg, "halt")) {
                                mode = UNINSTALL;
                                immediately = true;
                        } else if (!strcmp(arg, "info")) {
                                mode = INFO;
                        } else if (!strcmp(arg, "help")) {
                                print_usage();
                                return 0;
                        } else {
                                fprintf(stderr, "Invalid option `-%s'.\n", arg);
                                print_usage(true);
                                return 1;
                        }
                } else {
                        // "-" options
                        arg++;
                        int32 count = strlen(arg);
                        for (int i = 0; error == B_OK && i < count; i++) {
                                switch (arg[i]) {
                                        case 'h':
                                                print_usage();
                                                return 0;
                                        default:
                                                fprintf(stderr, "Invalid option `-%c'.\n", arg[i]);
                                                print_usage(true);
                                                return 1;
                                }
                        }
                }
        }

        // parse rest (the file name)
        if (argIndex != argc - 1) {
                print_usage(true);
                return 1;
        }
        const char *file = argv[argIndex];

        // do the job
        switch (mode) {
                case INSTALL:
                        error = install_file(file, size);
                        break;
                case UNINSTALL:
                        error = uninstall_file(file, immediately);
                        break;
                case INFO:
                        error = info_for_device(file);
                        break;
        }

        return (error == B_OK ? 0 : 1);
}