/*-
 * Copyright (c) 2024 Kyle Evans <kevans@FreeBSD.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <sys/param.h>

#include <err.h>
#include <errno.h>
#include <limits.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

static void
usage(void)
{

        fprintf(stderr, "usage: %s [-b bytes | -c lines | -e] [-s buffer-size]\n",
            getprogname());
        exit(1);
}

int
main(int argc, char *argv[])
{
        char *buf;
        const char *errstr;
        size_t bufsz = 0, reps;
        ssize_t ret;
        enum { MODE_BYTES, MODE_COUNT, MODE_EOF } mode;
        int ch;

        /*
         * -b specifies number of bytes.
         * -c specifies number of read() calls.
         * -e specifies eof (default)
         * -s to pass a buffer size
         *
         * Reading N lines is the same as -c with a high buffer size.
         */
        mode = MODE_EOF;
        while ((ch = getopt(argc, argv, "b:c:es:")) != -1) {
                switch (ch) {
                case 'b':
                        mode = MODE_BYTES;
                        reps = strtonum(optarg, 0, SSIZE_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "strtonum: %s", errstr);
                        break;
                case 'c':
                        mode = MODE_COUNT;
                        reps = strtonum(optarg, 1, SSIZE_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "strtonum: %s", errstr);
                        break;
                case 'e':
                        mode = MODE_EOF;
                        break;
                case 's':
                        bufsz = strtonum(optarg, 1, SSIZE_MAX, &errstr);
                        if (errstr != NULL)
                                errx(1, "strtonum: %s", errstr);
                        break;
                default:
                        usage();
                }
        }

        if (bufsz == 0) {
                if (mode == MODE_BYTES)
                        bufsz = reps;
                else
                        bufsz = LINE_MAX;
        }

        buf = malloc(bufsz);
        if (buf == NULL)
                err(1, "malloc");

        for (;;) {
                size_t readsz;

                /*
                 * Be careful not to over-read if we're in byte-mode.  In every other
                 * mode, we'll read as much as we can.
                 */
                if (mode == MODE_BYTES)
                        readsz = MIN(bufsz, reps);
                else
                        readsz = bufsz;

                ret = read(STDIN_FILENO, buf, readsz);
                if (ret == -1 && errno == EINTR)
                        continue;
                if (ret == -1)
                        err(1, "read");
                if (ret == 0) {
                        if (mode == MODE_EOF)
                                return (0);
                        errx(1, "premature EOF");
                }

                /* Write out what we've got */
                write(STDOUT_FILENO, buf, ret);

                /*
                 * Bail out if we've hit our metric (byte mode / count mode).
                 */
                switch (mode) {
                case MODE_BYTES:
                        reps -= ret;
                        if (reps == 0)
                                return (0);
                        break;
                case MODE_COUNT:
                        reps--;
                        if (reps == 0)
                                return (0);
                        break;
                default:
                        break;
                }
        }

        return (0);
}