/*
 * Copyright (c) 2019 Adélie Userland Team. All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 * 
 * 1. Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 * 
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 * 
 * 3. Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
 * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#include <sys/statvfs.h>

#include <bsd/stdlib.h>

#include <errno.h>
#include <locale.h>
#include <math.h>
#include <mntent.h>
#include <stdbool.h>
#include <stdio.h>
#define _POSIX_C_SOURCE 200809L
#include <string.h>
#include <unistd.h>

static void usage(void);
char *humanize(uint64_t size, bool si);

static bool hflag, kflag, Pflag, tflag, sflag, restricted;

struct fs {
    char *device;
    char *mount;
    char *type;
    char *total;
    char *used;
    char *avail;
    char *percent;
    unsigned long fsid;
    bool print;
    struct fs *next;
};

struct fs *fs_create(void) {
    struct fs *node = malloc(sizeof(struct fs));
    node->next = NULL;
    return node;
}

struct fs *fs_add(struct fs *root, const char *device, const char *mount,
                  const char *type, const char total[21], const char used[21],
                  const char avail[21], double percent, unsigned long fsid) {
    struct fs *node = fs_create();
    char temp[5];

    node->device = strndup(device, 24);
    node->total  = strndup(total, 20);
    node->used   = strndup(used, 20);
    node->avail  = strndup(avail, 20);
    node->mount  = strndup(mount, 64);
    node->type   = strndup(type, 10);
    snprintf(temp, 5, "%.0f%%", percent);
    node->percent = strndup(temp, 5);
    node->fsid = fsid;
    node->print = false;

    if (root == NULL) {
        return node;
    } else if (root->next == NULL) {
        root->next = node;
        return root;
    } else {
        struct fs *temp = root;
        while (temp != NULL) {
            if (temp->next == NULL) {
                temp->next = node;
                return root;
            }
            temp = temp->next;
        }
    }
}

int8_t longest[5];
#define R_FSNAME 0
#define R_TOTAL  1
#define R_USED   2
#define R_AVAIL  3
#define R_FSTYPE 4
#define R_MAX    5
struct fs *fs_build(struct fs *root, int unit) {
    struct statvfs stv;
    FILE *mnt;
    struct mntent *ment;

    for (uint_fast8_t i = 0; i < R_MAX; i++) {
        longest[i] = -1;
    }

    mnt = setmntent("/etc/mtab", "r");
    while ((ment = getmntent(mnt)) != NULL) {
        if (statvfs(ment->mnt_dir, &stv) != 0) {
            char *error = strerror(errno);
            fprintf(stderr, "statvfs(): unable to read %s: %s\n",
                    ment->mnt_dir, error);
            continue;
        }
        if (stv.f_blocks == 0) {
            continue;
        }
        int8_t current[5], spacelen = -1;
        char totals[21], useds[21], avails[21];
        for (uint_fast8_t i = 0; i < R_MAX; i++) {
            current[i] = -1;
        }
        uint64_t used = (stv.f_blocks - stv.f_bfree) * stv.f_frsize;
        uint64_t total = stv.f_blocks * stv.f_frsize;
        uint64_t avail = stv.f_bavail * stv.f_frsize;
        double percent = 100 * used / (double)total;

        if (hflag || sflag) {
            spacelen = 8;
            snprintf(totals, 8, "%s", humanize(total, sflag));
            snprintf(useds,  8, "%s", humanize(used,  sflag));
            snprintf(avails, 8, "%s", humanize(avail, sflag));
        } else {
            spacelen = 21;
            snprintf(totals, 21, "%ld", total / unit);
            snprintf(useds,  21, "%ld", used  / unit);
            snprintf(avails, 21, "%ld", avail / unit);
        }
        current[R_FSNAME] = strnlen(ment->mnt_fsname, 24);
        current[R_FSTYPE] = strnlen(ment->mnt_type, 10);
        current[R_TOTAL]  = strnlen(totals, spacelen);
        current[R_USED]   = strnlen(useds, spacelen);
        current[R_AVAIL]  = strnlen(avails, spacelen);

        for (uint_fast8_t i = 0; i < R_MAX; i++) {
            if (longest[i] < current[i]) {
                longest[i] = current[i];
            }
        }

        root = fs_add(root, ment->mnt_fsname, ment->mnt_dir, ment->mnt_type,
                      totals, useds, avails, percent, stv.f_fsid);
    }
    endmntent(mnt);

    return root;
}
void fs_filter(struct fs *root, int argc, char **argv) {
    struct statvfs stv;
    for (uint_fast16_t i = 0; i < argc; i++) {
        struct fs *temp = root;
        if (statvfs(argv[i], &stv) != 0) {
            char *error = strerror(errno);
            fprintf(stderr, "statvfs(): unable to read %s: %s\n",
                    argv[i], error);
            continue;
        }
        while (temp != NULL) {
            if (temp->fsid == stv.f_fsid) {
                temp->print = true;
                break;
            }
            temp = temp->next;
        }
    }
}

int
main(int argc, char **argv) {
    int ch, unit = 512;
    struct fs *root = NULL;
    hflag = kflag = Pflag = tflag = sflag = restricted = false;
    char *header = "Filesystem %d-blocks Used Avail Capacity Mounted on\n";
    char *fmt = "%s %s %s %s %s %s\n";
    int buflen = -1;

    setprogname(argv[0]);
    (void)setlocale(LC_ALL, "");

    while ((ch = getopt(argc, argv, "hksPt")) != -1) {
        switch(ch) {
        case 'h':
            if (kflag || sflag) {
                fprintf(stderr,
                        "%s: -h, -s, and -k cannot be used together.\n",
                        getprogname());
                usage();
            }
            hflag = true;
            unit = -1;
            break;
        case 'k':
            if (hflag || sflag) {
                fprintf(stderr,
                        "%s: -h, -s, and -k cannot be used together.\n",
                        getprogname());
                usage();
            }
            kflag = true;
            unit = 1024;
            break;
        case 'P':
            if (tflag) {
                fprintf(stderr, "%s: Cannot output types in POSIX mode.\n",
                        getprogname());
                usage();
            }
            Pflag = true;
            break;
        case 's':
            if (kflag || hflag) {
                fprintf(stderr,
                        "%s: -h, -s, and -k cannot be used together.\n",
                        getprogname());
                usage();
            }
            sflag = true;
            unit = -1;
            break;
        case 't':
            if (Pflag) {
                fprintf(stderr, "%s: Cannot output types in POSIX mode.\n",
                        getprogname());
                usage();
            }
            tflag = true;
            break;
        case '?':
        default:
            usage();
        }
    }
    argc -= optind;
    argv += optind;

    if (argv[0] != NULL) {
        /* only get information on filesystems containing the listed files */
        restricted = true;
    }

    root = fs_build(root, unit);

    if (restricted) {
        printf("restrict\n");
        fs_filter(root, argc, argv);
    }
    if (!Pflag) {
        for (uint_fast8_t i = 0; i < R_MAX; i++) {
            buflen += longest[i] + 1;
        }
        if (!tflag) {
            buflen -= longest[R_FSTYPE] + 1;
        }
    } else {
        buflen = strlen(fmt);
    }
    char temp[buflen + 1];
    if (!Pflag) {
        if (tflag) {
            snprintf(temp, buflen,
                     "%%-%ds %%%ds %%%ds %%%ds  %%4s  %%-%ds %%s\n",
                     longest[R_FSNAME], longest[R_TOTAL], longest[R_USED],
                     longest[R_AVAIL],  longest[R_FSTYPE]);
        } else {
            snprintf(temp, buflen,
                     "%%-%ds %%%ds %%%ds %%%ds  %%4s  %%s\n",
                     longest[R_FSNAME], longest[R_TOTAL], longest[R_USED],
                     longest[R_AVAIL]);
        }
    } else {
        snprintf(temp, buflen, "%s", fmt);
    }
    if (Pflag) {
        printf("Filesystem %d-blocks Used Avail Capacity Mounted on\n", unit);
    } else if (tflag) {
        printf(temp, "Filesystem", "Total", "Used", "Avail", "Use%", "Type",
               "Mounted on");
    } else {
        printf(temp, "Filesystem", "Total", "Used", "Avail", "Use%",
               "Mounted on");
    }
    while (root != NULL) {
        if ((restricted && root->print) || !restricted) {
            if (Pflag) {
                printf("%s %s %s %s %s %s\n", root->device, root->total, root->used, root->avail,
                       root->percent, root->mount);
            } else if (tflag) {
                printf(temp, root->device, root->total, root->used, root->avail,
                       root->percent, root->type, root->mount);
            } else {
                printf(temp, root->device, root->total, root->used, root->avail,
                       root->percent, root->mount);
            }
        }
        root = root->next;
    }
    return EXIT_SUCCESS;
}

char *scale[] = {
    "",
    "k",
    "M",
    "G",
    "T",
    "P",
    "E",
    "Z",
    "Y"
};
const uint8_t scale_end = 9;

char *humanize(uint64_t size, bool si) {
    uint16_t unit = si ? 1000 : 1024;
    int8_t radix = (int8_t)(log(size)/log(unit));
    char suffix[2], ret[10];
    double val;
    uint64_t temp = size;

    if (radix > 9) {
        radix = 9;
    }

    val = size / (double)pow(unit, radix);

    snprintf(ret, 10, "%.0f%s%sB", val, scale[radix], (si || (radix == 0)) ? "" : "i");

    return strndup(ret, 10);
}

void usage(void) {
    fprintf(stderr,
        "%s: [-hkPst] [file]...\n\n"
        "-h  Traditional human-readable units (1KiB is 1024 bytes); not usable with -k -P -s\n"
        "-k  1024-byte units; not usable with -h -s\n"
        "-P  POSIX-compliant output; not usable with -h -s -t\n"
        "-s  SI human-readable units (1KB is 1000 bytes); not usable with -h -k -P\n"
        "-t  Show filesystem type; not usable with -P\n", getprogname());
    exit(EXIT_FAILURE);
}