/*
 * Copyright (c) 2023 Benjamin Stürz <benni@stuerz.xyz>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */
#define _XOPEN_SOURCE 700
#define _BSD_SOURCE 1
#define DKTYPENAMES
#define WSDEBUG 0
#include <stddef.h>
#include <dev/biovar.h>
#include <sys/cdefs.h>
#include <sys/types.h>
#include <sys/disklabel.h>
#include <sys/sysctl.h>
#include <sys/ioctl.h>
#include <sys/mount.h>
#include <sys/dkio.h>
#include <inttypes.h>
#include <stdbool.h>
#include <string.h>
#include <stdarg.h>
#include <stdint.h>
#include <stdlib.h>
#include <unistd.h>
#include <libgen.h>
#include <stdio.h>
#include <fcntl.h>
#include <ctype.h>
#include <errno.h>
#include <util.h>
#include <err.h>

#define PREFIX0_ASCII   "|-"
#define PREFIX0_UNICODE "├─"
#define PREFIX1_ASCII   "`-"
#define PREFIX1_UNICODE "└─"
#define PREFIX2_ASCII   "| `-"
#define PREFIX2_UNICODE "│ └─"

static const char *prefix0, *prefix1, *prefix2;

static int diskcount (void)
{
    const int mib[2] = { CTL_HW, HW_DISKCOUNT };
    int diskcount;
    size_t len = sizeof diskcount;

    if (sysctl (mib, 2, &diskcount, &len, NULL, 0) == -1)
        err (1, "sysctl(hw.diskcount)");

    return diskcount;
}

static char *disknames (void)
{
    const int num = diskcount ();
    const int mib[2] = { CTL_HW, HW_DISKNAMES };
    size_t len = 32 * num;
    char *buffer = malloc (len);

    if (sysctl (mib, 2, buffer, &len, NULL, 0) == -1)
        err (1, "sysctl(hw.disknames)");

    return buffer;
}

static char *stripdisk (char *n)
{
    const char sufx[] = " disk";
    const size_t ln = strnlen (n, 16);
    const size_t ls = sizeof sufx - 1;

    if (memcmp (n + ln - ls, sufx, ls) == 0) {
        n[ln - ls] = '\0';
    }

    return n;
}

static void print_size (uint64_t sz)
{
    const struct unit {
        char sym;
        uint64_t factor;
    } units[] = {
        { 'P', 1ull << 50 },
        { 'T', 1ull << 40 },
        { 'G', 1ull << 30 },
        { 'M', 1ull << 20 },
        { 'K', 1ull << 10 },
        { '0', 0 },
    };

    char sym = 'B';
    uint64_t factor = 1;

    for (const struct unit *u = &units[0]; u->factor; ++u) {
        if (sz >= (u->factor * 9 / 10)) {
            sym = u->sym;
            factor = u->factor;
            break;
        }
    }

    const unsigned scaled10 = sz * 10 / factor;
    const unsigned scaled = sz / factor;
    if (scaled10 >= 1000) {
        printf ("%u", scaled);
    } else if (scaled10 >= 100) {
        printf (" %u", scaled);
    } else {
        printf ("%u.%u", scaled, scaled10 % 10);
    }

    putchar (sym);
    putchar (' ');
}

enum {
    FIELD_NAME      = 0x01,
    FIELD_DUID      = 0x02,
    FIELD_SIZE      = 0x04,
    FIELD_USED      = 0x08,
    FIELD_FREE      = 0x10,
    FIELD_TYPE      = 0x20,
    FIELD_COMMENT   = 0x40,

    FIELD_DEFAULT   = FIELD_NAME | FIELD_SIZE | FIELD_TYPE | FIELD_COMMENT,
};

enum {
    OPT_NOHEADER    = 0x01,
    OPT_NOBIO       = 0x02,
};

struct my_diskinfo;

struct my_partinfo {
    char letter;
    uint64_t size;
    uint64_t fssize;
    uint64_t free;
    const char *fstype;
    const char *mount;
    const struct my_diskinfo *sub; // If this is part of a RAID
    const char *raidstatus; // Only available if (sub != NULL)
};

struct my_diskinfo {
    char type[16];
    char label[16];
    char name[8];
    uint64_t size;
    uint64_t used;
    u_char duid[8];
    uint8_t num_parts;
    struct my_partinfo parts[MAXPARTITIONS];
    const char *raidstatus; // If this is a RAID device
};

struct padding {
    int name;
    /* duid = 8 */
    /* size = 4 */
    /* used = 4 */
    /* free = 4 */
    int type;
    int comment;
};

static void print_header (int fields, const struct padding *p)
{
    if (fields & FIELD_NAME)
        printf ("%-*s ", p->name, "NAME");

    if (fields & FIELD_DUID)
        printf ("%-18s ", "DUID");

    if (fields & FIELD_SIZE)
        printf ("%-4s ", "SIZE");

    if (fields & FIELD_USED)
        printf ("%-4s ", "USED");

    if (fields & FIELD_FREE)
        printf ("%-4s ", "FREE");

    if (fields & FIELD_TYPE)
        printf ("%-*s ", p->type, "TYPE");

    if (fields & FIELD_COMMENT)
        printf ("%-*s ", p->comment, "COMMENT");

#if WSDEBUG
    putchar ('X');
#endif

    putchar ('\n');
}

static void print_duid (const u_char *duid)
{
    for (size_t i = 0; i < 8; ++i) {
        printf ("%02x", duid[i]);
    }
}

static void print_disk (const struct my_diskinfo *, int, int, const char *, const struct padding *);
static void print_part (
    const struct my_diskinfo *disk,
    const struct my_partinfo *part,
    int fields,
    int options,
    bool last,
    const struct padding *p
) {
    if (fields & FIELD_NAME) {
        const char *prefix = last ? prefix1 : prefix0;
        printf (
            "%s%s%c%-*s ",
            prefix,
            disk->name,
            part->letter,
            p->name - 2 - (int)strlen (disk->name) - 1,
            ""
        );
    }

    if (fields & FIELD_DUID) {
        print_duid (disk->duid);
        printf (".%c ", part->letter);
    }

    if (fields & FIELD_SIZE)
        print_size (part->size);

    if (fields & FIELD_USED) {
        if (part->fssize) {
            print_size (part->fssize - part->free);
        } else {
            printf (" N/A ");
        }
    }

    if (fields & FIELD_FREE) {
        if (part->fssize) {
            print_size (part->free);
        } else {
            printf (" N/A ");
        }
    }

    if (fields & FIELD_TYPE)
        printf ("%-*s ", p->type, part->fstype);

    if (fields & FIELD_COMMENT)
        printf ("%-*s ", p->comment, part->mount ? part->mount : "");

#if WSDEBUG
    putchar ('X');
#endif

    putchar ('\n');

    if (part->sub) {
        print_disk (part->sub, fields, options, part->raidstatus, p);
    }
}

static void print_disk (
    const struct my_diskinfo *disk,
    int fields,
    int options,
    const char *raidstatus,
    const struct padding *p
) {
    if (fields & FIELD_NAME) {
        const char *prefix = raidstatus ? prefix2 : "";

        printf (
            "%s%-*s ",
            prefix,
            p->name - (raidstatus ? 4 : 0),
            disk->name
        );
    }

    if (fields & FIELD_DUID) {
        print_duid (disk->duid);
        printf ("   ");
    }

    if (fields & FIELD_SIZE)
        print_size (disk->size);

    if (fields & FIELD_USED)
        print_size (disk->used);

    if (fields & FIELD_FREE)
        print_size (disk->size - disk->used);

    if (fields & FIELD_TYPE)
        printf ("%-*.16s ", p->type, disk->type);

    if (fields & FIELD_COMMENT) {
        if (raidstatus) {
            printf ("%-*s ", p->comment, raidstatus);
        } else if (disk->raidstatus) {
            printf (
                "%.16s (%s)%*s ",
                disk->label,
                disk->raidstatus,
                p->comment - (int)strnlen (disk->label, sizeof disk->label) - (int)strlen (disk->raidstatus) - 3,
                ""
            );
        } else {
            printf ("%-*.16s ", p->comment, disk->label);
        }
    }

#if WSDEBUG
    putchar ('X');
#endif

    putchar ('\n');

    if (!raidstatus) {
        for (uint8_t i = 0; i < disk->num_parts; ++i)
            print_part (disk, &disk->parts[i], fields, options, i == (disk->num_parts - 1), p);
    }
}

static const struct statfs *find_mount (const char *dev)
{
    static struct statfs *mounts = NULL;
    static int n_mounts;

    if (!mounts) {
        n_mounts = getmntinfo (&mounts, MNT_NOWAIT);
        if (n_mounts == 0)
            err (1, "getmntinfo()");
    }

    for (int i = 0; i < n_mounts; ++i) {
        if (!strcmp (dev, mounts[i].f_mntfromname))
            return &mounts[i];
    }

    return NULL;
}

static int read_disk (const char *name, struct my_diskinfo *disk)
{
    struct disklabel label;
    char *ppath, *letter;

	memset (disk, 0, sizeof *disk);

    {   // Read disklabel.
        size_t len;
        int fd;

        fd = opendev (name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, &ppath);
        if (fd < 0) {
            warn ("read_disk(): opendev(%s)", name);
			return -1;
		}

        if (ioctl (fd, DIOCGDINFO, &label) < 0) {
            warn ("read_disk(): ioctl(%s, DIOCGDINFO)", name);
			close (fd);
			return -1;
		}
        close (fd);

        len = strlen (ppath);
        letter = ppath + len - 1;
    }

    strlcpy (disk->name, name, sizeof disk->name);
    disk->size = DL_GETDSIZE (&label) * label.d_secsize;
    memcpy (disk->type, label.d_typename, sizeof disk->type);
    stripdisk (disk->type);
    memcpy (disk->label, label.d_packname, sizeof disk->label);
    memcpy (disk->duid, label.d_uid, sizeof disk->duid);
    disk->num_parts = 0;

    for (uint16_t i = 0; i < label.d_npartitions; ++i) {
        const struct partition *p = &label.d_partitions[i];
        const struct statfs *mnt;
        struct my_partinfo part;

		memset (&part, 0, sizeof part);

        part.size = DL_GETPSIZE (p) * label.d_secsize;

        if (!part.size)
            continue;

        part.letter = 'a' + i;
        *letter = part.letter;
        mnt = find_mount (ppath);

        if (mnt) {
            const uint64_t bs = mnt->f_bsize;
            part.mount = mnt->f_mntonname;
            part.fssize = mnt->f_blocks * bs;
            part.free = mnt->f_bfree * bs;
        }

        part.fstype = fstypenames[p->p_fstype];
        if (i != 2)
            disk->used += part.size;
        disk->parts[disk->num_parts++] = part;
    }

    return 0;
}

static const char *bd_statusstr (int status) {
    switch (status) {
    case BIOC_SDONLINE:     return BIOC_SDONLINE_S;
    case BIOC_SDOFFLINE:    return BIOC_SDOFFLINE_S;
    case BIOC_SDFAILED:     return BIOC_SDFAILED_S;
    case BIOC_SDREBUILD:    return BIOC_SDREBUILD_S;
    case BIOC_SDHOTSPARE:   return BIOC_SDHOTSPARE_S;
    case BIOC_SDUNUSED:     return BIOC_SDUNUSED_S;
    case BIOC_SDSCRUB:      return BIOC_SDSCRUB_S;
    case BIOC_SDINVALID:    return BIOC_SDINVALID_S;
    default:                return "Unknown";
    }
}

static const char *bv_statusstr (int status) {
    switch (status) {
    case BIOC_SVONLINE:     return BIOC_SVONLINE_S;
    case BIOC_SVOFFLINE:    return BIOC_SVOFFLINE_S;
    case BIOC_SVDEGRADED:   return BIOC_SVDEGRADED_S;
    case BIOC_SVBUILDING:   return BIOC_SVBUILDING_S;
    case BIOC_SVSCRUB:      return BIOC_SVSCRUB_S;
    case BIOC_SVREBUILD:    return BIOC_SVREBUILD_S;
    case BIOC_SVINVALID:    return BIOC_SVINVALID_S;
    default:                return "Unknown";
    }
}

static void read_raid (
    struct my_diskinfo *disk,
    struct my_diskinfo *disks,
    size_t num_disks
) {
    struct bioc_inq bi;
    int fd;

    fd = opendev (disk->name, O_RDONLY, OPENDEV_PART | OPENDEV_BLCK, NULL);
    if (fd < 0) {
        warn ("read_raid(): opendev(%s)", disk->name);
        return;
    }

	memset (&bi, 0, sizeof bi);

    if (ioctl (fd, BIOCINQ, &bi) == -1)
        goto ret;

    for (int i = 0; i < bi.bi_novol; ++i) {
        struct bioc_vol bv;

		memset (&bv, 0, sizeof bv);
        memcpy (&bv.bv_bio, &bi.bi_bio, sizeof bi.bi_bio);
        bv.bv_volid = i;

        if (ioctl (fd, BIOCVOL, &bv) == -1) {
            warn ("read_raid(%s): BIOCVOL(%d)", disk->name, i);
            continue;
        }

        if (strcmp (disk->name, bv.bv_dev) != 0)
            continue;

        disk->raidstatus = bv_statusstr (bv.bv_status);

        for (int j = 0; j < bv.bv_nodisk; ++j) {
            struct bioc_disk bd;
            size_t len_vendor;
            char letter;

			memset (&bd, 0, sizeof bd);
            memcpy (&bd.bd_bio, &bi.bi_bio, sizeof bi.bi_bio);
            bd.bd_volid = i;
            bd.bd_diskid = j;

            if (ioctl (fd, BIOCDISK, &bd) == -1) {
                warn ("read_raid(%s): BIOCDISK(%d, %d)", disk->name, i, j);
                continue;
            }

            len_vendor = strlen (bd.bd_vendor);
            if (len_vendor < 4 || len_vendor > 8) {
                warnx ("read_raid(%s): unexpected vendor string: %.32s", disk->name, bd.bd_vendor);
                continue;
            }
            letter = bd.bd_vendor[len_vendor - 1];

            for (size_t k = 0; k < num_disks; ++k) {
                if (!memcmp (bd.bd_vendor, disks[k].name, 3)) {
                    for (size_t l = 0; l < disks[k].num_parts; ++l) {
                        if (letter == disks[k].parts[l].letter) {
                            disks[k].parts[l].sub = disk;
                            disks[k].parts[l].raidstatus = bd_statusstr(bd.bd_status);
                            goto found;
                        }
                    }
                }
            }
        found:;
        }
    }
ret:
    close (fd);
}

static int usage (void)
{
    fputs ("Usage: lsblk [-abinUuV] [disk...]\n", stderr);
    return 1;
}

static void pad_update (int *pad, int newval)
{
    if (newval > *pad)
        *pad = newval;
}

static void pad_disk (struct padding *p, const struct my_diskinfo *disk)
{
    size_t len_disk;
    int comment;

    len_disk = strnlen (disk->name, sizeof disk->name);
    comment = strnlen (disk->label, sizeof disk->label);
    if (disk->raidstatus)
        comment += strlen (disk->raidstatus) + 3;

    pad_update (&p->name, len_disk);
    pad_update (&p->type, strnlen (disk->type, sizeof disk->type));
    pad_update (&p->comment, comment);

    for (int i = 0; i < disk->num_parts; ++i) {
        const struct my_partinfo *part = &disk->parts[i];

        pad_update (&p->name, len_disk + 3);
        pad_update (&p->type, strlen (part->fstype));

        if (part->sub) {
            pad_update (&p->name, strnlen (part->sub->name, sizeof part->sub->name) + 4);
            pad_update (&p->comment, strlen (part->raidstatus));
        } else if (part->mount) {
            pad_update (&p->comment, strlen (part->mount));
        }
    }
}


static int compare_disk (const void *p1, const void *p2)
{
    const struct my_diskinfo *d1 = p1;
    const struct my_diskinfo *d2 = p2;

    return strcmp (d1->name, d2->name);
}

int main (int argc, char *argv[])
{
    int option;
    int fields = FIELD_DEFAULT;
    int options = 0;
	int ret = 0;
    bool ascii = false;

    if (unveil ("/dev", "r") == -1)
        err (1, "unveil(/dev)");

    if (unveil (NULL, NULL) == -1)
        err (1, "unveil()");

    while ((option = getopt (argc, argv, ":abinUuV")) != -1) {
        switch (option) {
        case 'a':
            fields = -1;
            break;
        case 'b':
            options |= OPT_NOBIO;
            break;
        case 'i':
            ascii = true;
            break;
        case 'n':
            options |= OPT_NOHEADER;
            break;
        case 'U':
            fields |= FIELD_DUID;
            break;
        case 'u':
            fields |= FIELD_USED | FIELD_FREE;
            break;
        case 'V':
            puts ("lsblk-" VERSION);
            return 0;
        default:
            return usage ();
        }
    }

    argv += optind;
    argc -= optind;

    char *names = argc == 0 ? disknames () : NULL;

    if (pledge ("stdio rpath disklabel", NULL) == -1)
        err (1, "pledge()");

    prefix0 = ascii ? PREFIX0_ASCII : PREFIX0_UNICODE;
    prefix1 = ascii ? PREFIX1_ASCII : PREFIX1_UNICODE;
    prefix2 = ascii ? PREFIX2_ASCII : PREFIX2_UNICODE;

    size_t cap_disks;
    if (argc == 0) {
        const char *s = names;
        cap_disks = 1;
        while ((s = strchr (s, ',')) != NULL)
            ++cap_disks, ++s;
    } else {
        cap_disks = argc;
    }

    struct my_diskinfo *disks = calloc (cap_disks, sizeof (struct my_diskinfo));
    size_t num_disks = 0;

    if (argc == 0) {
        for (char *name; (name = strsep (&names, ",")) != NULL; ) {
            char *colon = strchr (name, ':');
			struct my_diskinfo disk;

            if (colon)
                *colon = '\0';
			if (read_disk (name, &disk) == 0) {
				disks[num_disks++] = disk;
			} else {
				ret = 1;
			}

            if (colon)
                *colon = ':';
        }
    } else {
        for (int i = 0; i < argc; ++i) {
            char *name = basename (argv[i]);
            char *last = name + strlen (name) - 1;
			struct my_diskinfo disk;

            if (isalpha (*last)) {
                warnx ("%s: specifying a partition is not supported, using the drive.", name);
                *last = '\0';
            }
			if (read_disk (name, &disk) == 0) {
				disks[num_disks++] = disk;
			} else {
				ret = 1;
			}
        }
    }

    free (names);

	if (num_disks == 0)
		return ret;

    mergesort (disks, num_disks, sizeof *disks, compare_disk);

    if (!(options & OPT_NOBIO)) {
        for (size_t i = 0; i < num_disks; ++i) {
            read_raid (&disks[i], disks, num_disks);
        }
    }

    struct padding p = {
        .name = strlen ("NAME"),
        .type = strlen ("TYPE"),
        .comment = strlen ("COMMENT"),
    };

    for (size_t i = 0; i < num_disks; ++i)
        pad_disk (&p, &disks[i]);

    if (!(options & OPT_NOHEADER))
        print_header (fields, &p);

    for (size_t i = 0; i < num_disks; ++i) {
        print_disk (&disks[i], fields, options, NULL, &p);
    }

    return ret;
}

/* vim: set tabstop=4 shiftwidth=4 expandtab: */