/*
The MIT License (MIT)

Copyright © 2024 pacman64

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the “Software”), to deal
in the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
of the Software, and to permit persons to whom the Software is furnished to do
so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

/*
You can build this command-line app by running
    cc -Wall -s -O2 -o ./nh ./nh.c

Building with COMPACT_OUTPUT defined makes `nh` output many fewer bytes, at
the cost of using arguably worse colors.
*/

#include <fcntl.h>
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

// #define COMPACT_OUTPUT

// info is the multi-line help message
const char* info =
    ""
    "nh [options...] [filenames...]\n"
    "\n"
    "Nice Hexadecimal is a simple hexadecimal (base-16) viewer to inspect bytes\n"
    "from files or standard input.\n"
    "\n"
    "Each line shows the starting offset for the bytes shown, 16 of the bytes\n"
    "themselves in base-16 notation, and any ASCII codes when the byte values\n"
    "are in the typical ASCII range.\n"
    "\n"
    "The base-16 codes are color-coded, with most bytes shown in gray, while\n"
    "all-1 and all-0 bytes are shown in orange and blue respectively.\n"
    "\n"
    "All-0 bytes are the commonest kind in most binary file types and, along\n"
    "with all-1 bytes are also a special case worth noticing when exploring\n"
    "binary data, so it makes sense for them to stand out right away.\n"
    "\n"
    "\n"
    "Options\n"
    "\n"
    "    -h, --h            show this help message\n"
    "    -help, --help      aliases for option -h\n"
    "\n"
    "    -p, --p            plain-text output, without ANSI styles\n"
    "    -plain, --plain    aliases for option -p\n"
    "";

#ifdef COMPACT_OUTPUT
#define OUTPUT_FOR_00 "\x1b[34m00 "
#define OUTPUT_FOR_FF "\x1b[33mff "
#define NORMAL_HEX_STYLE "\x1b[37m"
#define ASCII_HEX_STYLE "\x1b[32m"
#define ASCII_BYTE_STYLE "\x1b[30m"
#else
#define OUTPUT_FOR_00 "\x1b[38;5;111m00 "
#define OUTPUT_FOR_FF "\x1b[38;5;209mff "
#define NORMAL_HEX_STYLE "\x1b[38;5;246m"
#define ASCII_HEX_STYLE "\x1b[38;5;72m"
#define ASCII_BYTE_STYLE "\x1b[38;5;239m"
#endif

// bufwriter is, as the name implies, a buffered-writer: when it's aimed at
// stdout, it considerably speeds up this app, as intended
typedef struct bufwriter {
    // buf is the buffer proper
    unsigned char* buf;

    // len is how many bytes of the buffer are currently being used
    size_t len;

    // cap is the capacity of the buffer, or the most bytes it can hold
    size_t cap;

    // out is the destination of all that's written into the buffer
    FILE* out;

    // done signals when/if no more output is accepted at the destination
    bool done;
} bufwriter;

// new_bufwriter is the constructor for type bufwriter
bufwriter new_bufwriter(FILE* dst, size_t cap) {
    bufwriter res;
    res.cap = cap;
    res.done = false;
    res.len = 0;
    res.out = dst;
    res.buf = malloc(res.cap);
    return res;
}

// flush does as it says: it empties the buffer after ensuring its bytes end
// on their intended destination
void flush(bufwriter* w) {
    if (w->len > 0 && fwrite(w->buf, w->len, 1, w->out) < 1) {
        w->done = true;
    }
    w->len = 0;
}

// close_bufwriter ensures all output is shown and deallocates the buffer
void close_bufwriter(bufwriter* w) {
    flush(w);
    free(w->buf);
    w->buf = NULL;
}

// write_bytes does as it says, minimizing the number of calls to fwrite
void write_bytes(bufwriter* w, const unsigned char* src, size_t len) {
    if (w->len + len < w->cap) {
        // all bytes fit into buffer
        memcpy(w->buf + w->len, src, len);
        w->len += len;
        return;
    }

    // ensure current buffer bytes go out, before crossing strides
    flush(w);

    // emit all chunks striding beyond/at the buffer's capacity
    for (; len >= w->cap; src += w->cap, len -= w->cap) {
        if (fwrite(src, w->cap, 1, w->out) < 1) {
            w->done = true;
            return;
        }
    }

    // now all, if any, remaining bytes will fit into the buffer
    memcpy(w->buf, src, len);
    w->len += len;
}

// write_byte does as it says
void write_byte(bufwriter* w, unsigned char b) {
    if (w->len >= w->cap) {
        flush(w);
    }

    unsigned char* ptr = w->buf + w->len;
    *ptr = b;
    w->len++;
}

// EMIT_CONST abstracts a common use-case of the bufwriter, which is
// emitting string constants without their final null byte
#define EMIT_CONST(w, x) write_bytes(w, (unsigned char*)x, sizeof(x) - 1)

// write_hex is faster than calling fprintf(w, "%02x", b): this matters
// because it's called for every input byte
void write_hex(bufwriter* w, unsigned char b) {
    const char* hex_digits = "0123456789abcdef";
    write_byte(w, hex_digits[b >> 4]);
    write_byte(w, hex_digits[b & 0x0f]);
}

// write_styled_hex emits an ANSI color-coded hexadecimal representation
// of the byte given
void write_styled_hex(bufwriter* w, unsigned char b) {
    // all-bits-off is almost always noteworthy
    if (b == 0) {
        EMIT_CONST(w, OUTPUT_FOR_00);
        return;
    }
    // all-bits-on is often noteworthy
    if (b == 0xff) {
        EMIT_CONST(w, OUTPUT_FOR_FF);
        return;
    }

    // regular ASCII display symbols
    if (32 <= b && b <= 126) {
        EMIT_CONST(w, ASCII_HEX_STYLE);
        write_hex(w, b);
        EMIT_CONST(w, ASCII_BYTE_STYLE);
        write_byte(w, b);
        return;
    }

    // ASCII control values, and other bytes beyond displayable ASCII
    EMIT_CONST(w, NORMAL_HEX_STYLE);
    write_hex(w, b);
    write_byte(w, ' ');
}

// ruler emits a ruler-like string of spaced-out symbols
void ruler(bufwriter* w, size_t bytes_per_line) {
    const size_t gap = 4;
    if (bytes_per_line < gap) {
        return;
    }

    EMIT_CONST(w, "             ·");
    for (size_t n = bytes_per_line - gap; n >= gap; n -= gap) {
        EMIT_CONST(w, "           ·");
    }
}

// write_commas_uint shows a number by separating 3-digits groups with commas
void write_commas_uint(bufwriter* w, size_t n) {
    if (n == 0) {
        EMIT_CONST(w, "0");
        return;
    }

    size_t digits;
    // 20 is the most digits unsigned 64-bit ints can ever need
    unsigned char buf[24];
    for (digits = 0; n > 0; digits++, n /= 10) {
        buf[sizeof(buf) - 1 - digits] = (n % 10) + '0';
    }

    // now emit the leading digits, which may not come in 3
    size_t leading = digits % 3;
    if (leading == 0) {
        // avoid having a comma before the first digit
        leading = digits < 3 ? digits : 3;
    }
    unsigned char* start = buf + sizeof(buf) - digits;
    write_bytes(w, start, leading);
    start += leading;
    digits -= leading;

    // now emit all remaining digits in groups of 3, alternating styles
    for (; digits > 0; start += 3, digits -= 3) {
        write_byte(w, ',');
        write_bytes(w, start, 3);
    }
}

// output_state ties all values representing the current state shared across
// all functions involved in interpreting the input-buffer and showing its
// bytes and ASCII values
typedef struct output_state {
    // the whole input-buffer and its currently-used length in bytes
    unsigned char* buf;
    size_t buflen;

    // the ASCII-text buffer and its currently-used length in bytes
    unsigned char* txt;
    size_t txtlen;

    // offset is the byte counter, shown at the start of each line
    size_t offset;

    // linewidth is how many bytes each line can show at most
    size_t linewidth;

    // lines is the line counter, which is used to provide periodic
    // breather lines, to make eye-scanning big output blobs easier
    size_t lines;

    // showtxt is a hint on whether it's sensible to show the ASCII-text
    // buffer for the current line
    bool showtxt;
} output_state;

// peek_ascii looks 2 lines ahead in the buffer to get all ASCII-like runs
// of bytes, which are later meant to show on the side panel
void peek_ascii(size_t i, size_t end, output_state* os) {
    unsigned char prev = 0;
    os->txtlen = 0;

    for (size_t j = i; j < end; j++) {
        const unsigned char b = os->buf[j];

        if (' ' < b && b <= '~') {
            bool first = os->txtlen == 0;
            if (first) {
                // show ASCII panel, if the symbols start on the current line
                os->showtxt = j - i < os->linewidth;
            }

            // add a space before the symbol, when it's the start of a `word`
            if ((prev <= ' ' || prev > '~') && !first) {
                os->txt[os->txtlen] = ' ';
                os->txtlen++;
            }

            // add the symbol itself
            os->txt[os->txtlen] = b;
            os->txtlen++;
        }

        prev = b;
    }
}

// write_plain_uint is the unstyled counterpart of func write_styled_uint
void write_plain_uint(bufwriter* w, size_t n) {
    if (n < 1) {
        EMIT_CONST(w, "       0");
        return;
    }

    size_t digits;
    // 20 is the most digits unsigned 64-bit ints can ever need
    unsigned char buf[24];
    for (digits = 0; n > 0; digits++, n /= 10) {
        buf[sizeof(buf) - 1 - digits] = (n % 10) + '0';
    }

    // left-pad the coming digits up to 8 chars
    if (digits < 8) {
        write_bytes(w, (unsigned char*)"        ", 8 - digits);
    }

    // emit all digits
    unsigned char* start = buf + sizeof(buf) - digits;
    write_bytes(w, start, digits);
}

// write_styled_uint is a quick way to emit the offset-counter showing at the
// start of each line; it assumes 8-item left-padding of values, unless the
// numbers are too big for that
void write_styled_uint(bufwriter* w, size_t n) {
    if (n < 1) {
        EMIT_CONST(w, "       0");
        return;
    }

    size_t digits;
    // 20 is the most digits unsigned 64-bit ints can ever need
    unsigned char buf[24];
    for (digits = 0; n > 0; digits++, n /= 10) {
        buf[sizeof(buf) - 1 - digits] = (n % 10) + '0';
    }

    // left-pad the coming digits up to 8 chars
    if (digits < 8) {
        write_bytes(w, (unsigned char*)"        ", 8 - digits);
    }

    // now emit the leading digits, which may be fewer than 3
    size_t leading = digits % 3;
    unsigned char* start = buf + sizeof(buf) - digits;
    write_bytes(w, start, leading);
    start += leading;
    digits -= leading;

    // now emit all remaining digits in groups of 3, alternating styles
    bool styled = leading != 0;
    for (; digits > 0; start += 3, digits -= 3, styled = !styled) {
        if (styled) {
            EMIT_CONST(w, "\x1b[38;5;243m");
            write_bytes(w, start, 3);
            EMIT_CONST(w, "\x1b[0m");
        } else {
            write_bytes(w, start, 3);
        }
    }
}

// emit_styled_file_info emits an ANSI-styled line showing a filename and the
// file's size in bytes
void emit_styled_file_info(bufwriter* w, const char* path, size_t nbytes) {
    EMIT_CONST(w, "• ");
    write_bytes(w, (unsigned char*)path, strlen(path));
    EMIT_CONST(w, "  \x1b[38;5;245m(");
    write_commas_uint(w, nbytes);
    EMIT_CONST(w, " bytes)\x1b[0m\n");
}

// emit_plain_file_info is the unstyled counterpart of func emit_styled_file_info
void emit_plain_file_info(bufwriter* w, const char* path, size_t nbytes) {
    EMIT_CONST(w, "• ");
    write_bytes(w, (unsigned char*)path, strlen(path));
    EMIT_CONST(w, "  (");
    write_commas_uint(w, nbytes);
    EMIT_CONST(w, " bytes)\n");
}

// emit_styled_line handles the details of showing a styled line out of the current
// input-buffer chunk
void emit_styled_line(bufwriter* w, size_t i, size_t end, output_state* os) {
    for (size_t j = i; j < end; j++, os->offset++) {
        const unsigned char b = os->buf[j];

        if (j % os->linewidth == 0) {
            // show a ruler every few lines to make eye-scanning easier
            if (os->lines % 5 == 0 && os->lines > 0) {
                EMIT_CONST(w, "        \x1b[38;5;245m");
                ruler(w, os->linewidth);
                EMIT_CONST(w, "\x1b[0m\n");
            }
            os->lines++;

            // start next line with offset of its 1st item, also
            // changing the background color for the colored hex
            // code which will follow
            // fprintf(stdout, "%8d", os->offset);
            write_styled_uint(w, os->offset);
            EMIT_CONST(w, "  \x1b[48;5;254m");
        }

        // show the current byte `with style`
        write_styled_hex(w, b);
    }

    if (os->showtxt) {
        EMIT_CONST(w, "\x1b[0m  ");
        for (size_t j = end - i; j < os->linewidth; j++) {
            EMIT_CONST(w, "   ");
        }

        write_bytes(w, os->txt, os->txtlen);
        write_byte(w, '\n');
        return;
    }
    EMIT_CONST(w, "\x1b[0m\n");
}

// emit_plain_line handles the details of showing a plain (unstyled) line out
// of the current input-buffer chunk
void emit_plain_line(bufwriter* w, size_t i, size_t end, output_state* os) {
    for (size_t j = i; j < end; j++, os->offset++) {
        const unsigned char b = os->buf[j];

        if (j % os->linewidth == 0) {
            // show a ruler every few lines to make eye-scanning easier
            if (os->lines % 5 == 0 && os->lines > 0) {
                // EMIT_CONST(w, "        ");
                // ruler(w, os->linewidth);
                write_byte(w, '\n');
            }
            os->lines++;

            // start next line with offset of its 1st item, also
            // changing the background color for the colored hex
            // code which will follow
            // fprintf(stdout, "%8d", os->offset);
            write_plain_uint(w, os->offset);
            EMIT_CONST(w, "  ");
        }

        // show the current byte `with style`
        write_hex(w, b);
        write_byte(w, ' ');
    }

    if (os->showtxt) {
        EMIT_CONST(w, "  ");
        for (size_t j = end - i; j < os->linewidth; j++) {
            EMIT_CONST(w, "   ");
        }

        write_bytes(w, os->txt, os->txtlen);
        write_byte(w, '\n');
        return;
    }
    write_byte(w, '\n');
}

// config has all the settings used to emit output
typedef struct config {
    // bytes_per_line determines the `width` of output lines
    size_t bytes_per_line;

    // emit_file_info is chosen to emit file-info with colors or plainly
    void (*emit_file_info)(bufwriter* w, const char* path, size_t nbytes);

    // emit_line is chosen to emit hex bytes with colors or plainly
    void (*emit_line)(bufwriter* w, size_t i, size_t end, output_state* os);
} config;

// handle_reader shows all bytes read from the source given as colored hex
// values, showing offsets and ASCII symbols on the sides of each output line
void handle_reader(bufwriter* w, FILE* src, config cfg) {
    const size_t bufcap = 32 * 1024;
    // limit line-width to the buffer's capacity
    if (cfg.bytes_per_line > bufcap) {
        cfg.bytes_per_line = bufcap;
    }

    const size_t two_lines = 2 * cfg.bytes_per_line;
    unsigned char txt[two_lines];

    unsigned char buf[bufcap];
    // ensure the effective buffer-size is a multiple of the line-width
    size_t max = bufcap - bufcap % cfg.bytes_per_line;

    output_state os;
    os.buf = buf;
    os.linewidth = cfg.bytes_per_line;
    os.lines = 0;
    os.offset = 0;
    os.txt = txt;

    const size_t one_line = cfg.bytes_per_line;

    while (!w->done) {
        os.buflen = fread(&buf, sizeof(unsigned char), max, src);
        if (os.buflen < 1) {
            // assume input is over when no bytes were read
            return;
        }

        for (size_t i = 0; i < os.buflen; i += one_line) {
            size_t end;

            // remember all ASCII symbols in current pair of output lines
            end = i + two_lines < os.buflen ? i + two_lines : os.buflen;
            peek_ascii(i, end, &os);

            // show current output line
            end = i + one_line < os.buflen ? i + one_line : os.buflen;
            cfg.emit_line(w, i, end, &os);
        }
    }
}

// handle_file handles data from the filename given; returns false only when
// the file can't be opened
bool handle_file(bufwriter* w, const char* path, config cfg) {
    // a `-` filename stands for the standard input
    if (strcmp(path, "-") == 0) {
        EMIT_CONST(w, "• <stdin>\n");
        EMIT_CONST(w, "\n");
        handle_reader(w, stdin, cfg);
        return true;
    }

    FILE* f = fopen(path, "rb");
    if (f == NULL) {
        // ensure currently-buffered/deferred output shows up right now: not
        // doing so may scramble results in the common case where stdout and
        // stderr are the same, thus confusing users
        flush(w);

        fprintf(stderr, "\x1b[31mcan't open file named %s\x1b[0m\n", path);
        return false;
    }

    // get the file size
    struct stat st;
    fstat(fileno(f), &st);

    // show output
    cfg.emit_file_info(w, path, st.st_size);
    EMIT_CONST(w, "\n");
    handle_reader(w, f, cfg);

    fclose(f);
    return true;
}

// is_help_option simplifies control-flow for func run
bool is_help_option(char* s) {
    return false ||
           strcmp(s, "-h") == 0 ||
           strcmp(s, "-help") == 0 ||
           strcmp(s, "--h") == 0 ||
           strcmp(s, "--help") == 0;
}

// is_plain_option simplifies control-flow for func run
bool is_plain_option(char* s) {
    return false ||
           strcmp(s, "-p") == 0 ||
           strcmp(s, "-plain") == 0 ||
           strcmp(s, "--p") == 0 ||
           strcmp(s, "--plain") == 0;
}

// run returns the number of errors
size_t run(int argc, char** argv) {
    config cfg;
    cfg.bytes_per_line = 16;
    cfg.emit_line = &emit_styled_line;
    cfg.emit_file_info = &emit_styled_file_info;

    // handle special cmd-line options and count filenames
    size_t fnames = 0;
    for (size_t i = 1; i < argc; i++) {
        if (is_help_option(argv[i])) {
            // help option is handled right away, also quitting the app
            fprintf(stderr, "%s", info);
            return 0;
        }
        if (is_plain_option(argv[i])) {
            cfg.emit_line = &emit_plain_line;
            cfg.emit_file_info = &emit_plain_file_info;
            continue;
        }
        fnames++;
    }

    bufwriter w = new_bufwriter(stdout, 32 * 1024);

    // no filenames means use stdin as the only input
    if (fnames == 0) {
        EMIT_CONST(&w, "• <stdin>\n");
        EMIT_CONST(&w, "\n");
        handle_reader(&w, stdin, cfg);
        close_bufwriter(&w);
        return 0;
    }

    size_t errors = 0;
    bool first_file = true;

    // handle all filenames given
    for (size_t i = 1; i < argc && !w.done; i++) {
        if (i == 1 && is_plain_option(argv[i])) {
            // special cmd-line options aren't filenames
            continue;
        }

        if (!first_file) {
            // put an empty line between adjacent hex outputs
            write_byte(&w, '\n');
        }

        if (!handle_file(&w, argv[i], cfg)) {
            errors++;
        }
        first_file = false;
    }

    close_bufwriter(&w);
    return errors;
}

int main(int argc, char** argv) {
#ifdef _WIN32
    setmode(fileno(stdin), O_BINARY);
    // ensure output lines end in LF instead of CRLF on windows
    setmode(fileno(stdout), O_BINARY);
    setmode(fileno(stderr), O_BINARY);
#endif

    // disable automatic stdio buffering, in favor of explicit buffering
    setvbuf(stdin, NULL, _IONBF, 0);
    setvbuf(stdout, NULL, _IONBF, 0);
    setvbuf(stderr, NULL, _IONBF, 0);

    return run(argc, argv) == 0 ? 0 : 1;
}