// Copyright 2017 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

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

#include <fuchsia/hardware/skipblock/c/fidl.h>
#include <lib/fdio/fd.h>
#include <lib/fdio/fdio.h>
#include <lib/fdio/directory.h>
#include <zircon/process.h>
#include <zircon/syscalls.h>

int usage(void) {
    fprintf(stderr, "usage: dd [OPTIONS]\n");
    fprintf(stderr, "dd can be used to convert and copy files\n");
    fprintf(stderr, " bs=BYTES  : read and write BYTES bytes at a time\n");
    fprintf(stderr, " count=N   : copy only N input blocks\n");
    fprintf(stderr, " ibs=BYTES : read BYTES bytes at a time (default: 512)\n");
    fprintf(stderr, " if=FILE   : read from FILE instead of stdin\n");
    fprintf(stderr, " obs=BYTES : write BYTES bytes at a time (default: 512)\n");
    fprintf(stderr, " of=FILE   : write to FILE instead of stdout\n");
    fprintf(stderr, " seek=N    : skip N obs-sized blocks at start of output\n");
    fprintf(stderr, " skip=N    : skip N ibs-sized blocks at start of input\n");
    fprintf(stderr, " N and BYTES may be followed by the following multiplicitive\n"
                    " suffixes: c = 1, w = 2, b = 512, kB = 1000, K = 1024,\n"
                    "           MB = 1000 * 1000, M = 1024 * 1024, xM = M,\n"
                    "           GB = 1000 * 1000 * 1000, G = 1024 * 1024 * 1024\n");
    fprintf(stderr, " --help : Show this help message\n");
    return -1;
}

// Returns "true" if the argument matches the prefix.
// In this case, moves the argument past the prefix.
bool prefix_match(const char** arg, const char* prefix) {
    if (!strncmp(*arg, prefix, strlen(prefix))) {
        *arg += strlen(prefix);
        return true;
    }
    return false;
}

#define MAYBE_MULTIPLY_SUFFIX(str, out, suffix, value) \
        if (!strcmp((str), (suffix))) {                \
            (out) *= (value);                          \
            return 0;                                  \
        }

// Parse the formatted size string from |s|, and place
// the result in |out|.
//
// Returns 0 on success.
int parse_size(const char* s, size_t* out) {
    char* endptr;
    if (!(*s >= '0' && *s <= '9')) {
        goto done;
    }
    *out = strtol(s, &endptr, 10);
    if (*endptr == '\0') {
        return 0;
    }

    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "G", 1UL << 30);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "GB", 1000 * 1000 * 1000UL);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "xM", 1UL << 20);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "M", 1UL << 20);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "MB", 1000 * 1000UL);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "K", 1UL << 10);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "kB", 1000UL);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "b", 512UL);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "w", 2UL);
    MAYBE_MULTIPLY_SUFFIX(endptr, *out, "c", 1UL);

done:
    fprintf(stderr, "Couldn't parse size string: %s\n", s);
    return -1;
}

typedef struct {
    bool use_count;
    size_t count;
    size_t input_bs;
    size_t input_skip;
    size_t output_bs;
    size_t output_seek;
    char input[PATH_MAX];
    char output[PATH_MAX];
} dd_options_t;

int parse_args(int argc, const char** argv, dd_options_t* options) {
    while (argc > 1) {
        const char* arg = argv[1];
        if (prefix_match(&arg, "bs=")) {
            size_t size;
            if (parse_size(arg, &size)) {
                return usage();
            }
            options->input_bs = size;
            options->output_bs = size;
        } else if (prefix_match(&arg, "count=")) {
            if (parse_size(arg, &options->count)) {
                return usage();
            }
            options->use_count = true;
        } else if (prefix_match(&arg, "ibs=")) {
            if (parse_size(arg, &options->input_bs)) {
                return usage();
            }
        } else if (prefix_match(&arg, "obs=")) {
            if (parse_size(arg, &options->output_bs)) {
                return usage();
            }
        } else if (prefix_match(&arg, "seek=")) {
            if (parse_size(arg, &options->output_seek)) {
                return usage();
            }
        } else if (prefix_match(&arg, "skip=")) {
            if (parse_size(arg, &options->input_skip)) {
                return usage();
            }
        } else if (prefix_match(&arg, "if=")) {
            strncpy(options->input, arg, PATH_MAX);
            options->input[PATH_MAX - 1] = '\0';
        } else if (prefix_match(&arg, "of=")) {
            strncpy(options->output, arg, PATH_MAX);
            options->output[PATH_MAX - 1] = '\0';
        } else {
            return usage();
        }
        argc--;
        argv++;
    }
    return 0;
}

#define MIN(x,y) ((x) < (y) ? (x) : (y))
#define MAX(x,y) ((x) < (y) ? (y) : (x))

bool is_skip_block(const char* path, fuchsia_hardware_skipblock_PartitionInfo* partition_info) {
    int fd = open(path, O_RDONLY);
    if (fd < 0) {
        return false;
    }

    zx_handle_t channel;
    zx_status_t status = fdio_get_service_handle(fd, &channel);
    if (status != ZX_OK) {
        return false;
    }

    // |status| is used for the status of the whole FIDL request. We expect
    // |status| to be ZX_OK if the channel connects to a skip-block driver.
    // |op_status| refers to the status of the underlying read/write operation
    // and will be ZX_OK only if the read/write succeeds. It is NOT set if
    // the channel is not connected to a skip-block driver.
    zx_status_t op_status;

    status = fuchsia_hardware_skipblock_SkipBlockGetPartitionInfo(
        channel, &op_status, partition_info);

    close(fd);
    return status == ZX_OK;
}

int main(int argc, const char** argv) {
    dd_options_t options;
    memset(&options, 0, sizeof(dd_options_t));
    options.input_bs = 512;
    options.output_bs = 512;
    int r;
    if ((r = parse_args(argc, argv, &options))) {
        return r;
    }

    if (options.input_bs == 0 || options.output_bs == 0) {
        fprintf(stderr, "block sizes must be greater than zero\n");
        return -1;
    }

    options.input_skip *= options.input_bs;
    options.output_seek *= options.output_bs;
    size_t buf_size = MAX(options.output_bs, options.input_bs);

    // Input and output fds
    int in = -1;
    int out = -1;
    // Buffer to contain partially read data
    char* buf = NULL;
    // Return code
    r = -1;
    // Number of full records copied to/from target
    size_t records_in = 0;
    size_t records_out = 0;
    // Size of remaining "partial" transfer from input / to output.
    size_t record_in_partial = 0;
    size_t record_out_partial = 0;
    // Logic for skip-block devices.
    fuchsia_hardware_skipblock_PartitionInfo partition_info = {};
    bool in_is_skip_block = false;
    bool out_is_skip_block = false;
    zx_handle_t vmo = ZX_HANDLE_INVALID;
    zx_handle_t channel_in = ZX_HANDLE_INVALID;
    zx_handle_t channel_out = ZX_HANDLE_INVALID;

    if (*options.input == '\0') {
        in = STDIN_FILENO;
    } else {
        in = open(options.input, O_RDONLY);
        if (in < 0) {
            fprintf(stderr, "Couldn't open input file %s : %d\n", options.input, errno);
            goto done;
        }
    }

    if (*options.output == '\0') {
        out = STDOUT_FILENO;
    } else {
        out = open(options.output, O_WRONLY | O_CREAT);
        if (out < 0) {
            fprintf(stderr, "Couldn't open output file %s : %d\n", options.output, errno);
            goto done;
        }
    }

    if (is_skip_block(options.input, &partition_info)) {
        if (options.input_bs % partition_info.block_size_bytes) {
            fprintf(stderr, "BS must be a multiple of %lu\n", partition_info.block_size_bytes);
            return false;
            goto done;
        }

        in_is_skip_block = true;
        fdio_get_service_handle(in, &channel_in);
    }

    if (is_skip_block(options.output, &partition_info)) {
        if (options.output_bs % partition_info.block_size_bytes) {
            fprintf(stderr, "BS must be a multiple of %lu\n", partition_info.block_size_bytes);
            return false;
            goto done;
        }

        out_is_skip_block = true;
        fdio_get_service_handle(out, &channel_out);
    }

    if (in_is_skip_block || out_is_skip_block) {
        if (zx_vmo_create(buf_size, 0, &vmo) != ZX_OK) {
            fprintf(stderr, "No memory\n");
            goto done;
        }
        if (zx_vmar_map(zx_vmar_root_self(), ZX_VM_PERM_READ | ZX_VM_PERM_WRITE,
                        0, vmo, 0, buf_size, (zx_vaddr_t*)&buf) != ZX_OK) {
            fprintf(stderr, "Failed to map vmo\n");
            goto done;
        }
    } else {
        buf = malloc(buf_size);
        if (buf == NULL) {
            fprintf(stderr, "No memory\n");
            goto done;
        }
    }

    if (options.input_skip != 0 && !in_is_skip_block) {
        // Try seeking first; if that doesn't work, try reading to an input buffer.
        if (lseek(in, options.input_skip, SEEK_SET) != (off_t) options.input_skip) {
            while (options.input_skip) {
                if (read(in, buf, options.input_bs) != (ssize_t) options.input_bs) {
                    fprintf(stderr, "Couldn't read from input\n");
                    goto done;
                }
                options.input_skip -= options.input_bs;
            }
        }
    }

    if (options.output_seek != 0 && !out_is_skip_block) {
        if (lseek(out, options.output_seek, SEEK_SET) != (off_t) options.output_seek) {
            fprintf(stderr, "Failed to seek on output\n");
            goto done;
        }
    }

    if (MAX(options.input_bs, options.output_bs) %
        MIN(options.input_bs, options.output_bs) != 0) {
        // TODO(smklein): Implement this case, rather than returning an error
        fprintf(stderr, "Input and output block sizes must be multiples\n");
        goto done;
    }

    bool terminating = false;
    size_t rlen = 0;
    while (true) {
        if (options.use_count && !options.count) {
            r = 0;
            goto done;
        }

        // Read as much as we can (up to input_bs) into our target buffer
        if (in_is_skip_block) {
            zx_handle_t dup;
            if (zx_handle_duplicate(vmo, ZX_RIGHT_SAME_RIGHTS, &dup) != ZX_OK) {
                fprintf(stderr, "Cannot duplicate handle\n");
                goto done;
            }
            const uint32_t block_count = (uint32_t)(options.input_bs /
                                                    partition_info.block_size_bytes);
            fuchsia_hardware_skipblock_ReadWriteOperation op = {
                .vmo = dup,
                .vmo_offset = 0,
                .block = (uint32_t)((options.input_skip / partition_info.block_size_bytes) +
                                    (records_in * block_count)),
                .block_count = block_count,
            };

            zx_status_t status;
            fuchsia_hardware_skipblock_SkipBlockRead(channel_in, &op, &status);

            if (status != ZX_OK) {
                fprintf(stderr, "Failed to read\n");
                goto done;
            }
            records_in++;
            rlen += options.input_bs;
        } else {
            ssize_t rout;
            if ((rout = read(in, buf, options.input_bs)) != (ssize_t) options.input_bs) {
                terminating = true;
            }
            if (rout == (ssize_t) options.input_bs) {
                records_in++;
            } else if (rout > 0) {
                record_in_partial = rout;
            }
            if (rout > 0) {
                rlen += rout;
            }
        }
        if (options.use_count) {
            --options.count;
            if (options.count == 0) {
                terminating = true;
            }
        }

        // If we can (or should, due to impending termination), dump the read
        // buffer into the output file.
        if (rlen >= options.output_bs || terminating) {
            if (out_is_skip_block) {
                zx_handle_t dup;
                if (zx_handle_duplicate(vmo, ZX_RIGHT_SAME_RIGHTS, &dup) != ZX_OK) {
                    fprintf(stderr, "Cannot duplicate handle\n");
                    goto done;
                }
                const uint32_t block_count = (uint32_t)(options.output_bs /
                                                        partition_info.block_size_bytes);

                fuchsia_hardware_skipblock_ReadWriteOperation op = {
                    .vmo = dup,
                    .vmo_offset = 0,
                    .block = (uint32_t)((options.output_seek / partition_info.block_size_bytes) +
                                        (records_out * block_count)),
                    .block_count = block_count,
                };

                zx_status_t status;
                bool bad_block_grown;
                fuchsia_hardware_skipblock_SkipBlockWrite(
                    channel_out, &op, &status, &bad_block_grown);

                if (status != ZX_OK) {
                    fprintf(stderr, "Failed to write\n");
                    goto done;
                }

                records_out++;
            } else {
                size_t off = 0;
                while (off != rlen) {
                    size_t wlen = MIN(options.output_bs, rlen - off);
                    if (write(out, buf + off, wlen) != (ssize_t) wlen) {
                        fprintf(stderr, "Couldn't write %zu bytes to output\n", wlen);
                        goto done;
                    }
                    if (wlen == options.output_bs) {
                        records_out++;
                    } else {
                        record_out_partial = wlen;
                    }
                    off += wlen;
                }
            }
            rlen = 0;
        }

        if (terminating) {
            r = 0;
            goto done;
        }
    }

done:
    printf("%zu+%u records in\n", records_in, record_in_partial ? 1 : 0);
    printf("%zu+%u records out\n", records_out, record_out_partial ? 1 : 0);
    printf("%zu bytes copied\n", records_out * options.output_bs + record_out_partial);

    if (buf) {
        if (in_is_skip_block || out_is_skip_block) {
            zx_vmar_unmap(zx_vmar_root_self(), (uintptr_t)*buf, buf_size);
        } else {
            free(buf);
        }
    }
    if (vmo != ZX_HANDLE_INVALID) {
        zx_handle_close(vmo);
    }

    if (in != -1) {
        close(in);
    }
    if (out != -1) {
        close(out);
    }
    return r;
}