#!/usr/bin/python3

# 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.


# Note: string slicing is a major source of inefficiencies in this script,
# making it viable only for small inputs; it's not clear what the stdlib
# offers to loop over sub-strings without copying data, which is really
# needed in this case.
#
# In the end the code has become much uglier by using explicit index-pairs,
# which are used/updated all over to avoid copying sub-strings. Standard
# output is already line-buffered by default, which is makes writing to it
# already fairly fast.


from io import TextIOWrapper
from sys import argv, exit, stderr, stdin, stdout


info = '''
nn [option...] [filepaths/URIs...]


Nice Numbers restyles all runs of 4+ digits by alternating ANSI-styles
every 3-digit group, so long numbers become easier to read at a glance.

All (optional) leading options start with either single or double-dash,
and most of them change the style/color used. Some of the options are,
shown in their single-dash form:

    -h          show this help message
    -help       show this help message

    -b          use a blue color
    -blue       use a blue color
    -bold       bold-style digits
    -g          use a green color
    -gray       use a gray color (default)
    -green      use a green color
    -hi         use a highlighting/inverse style
    -highlight  use a highlighting/inverse style
    -hilite     use a highlighting/inverse style
    -inverse    use a highlighting/inverse style
    -m          use a magenta color
    -magenta    use a magenta color
    -o          use an orange color
    -orange     use an orange color
    -p          use a purple color
    -purple     use a purple color
    -r          use a red color
    -red        use a red color
    -u          underline digits
    -underline  underline digits
'''

# handle standard help cmd-line options, quitting right away in that case
if len(argv) == 2 and argv[1] in ('-h', '--h', '-help', '--help'):
    print(info.strip(), file=stderr)
    exit(0)

# names_aliases normalizes lookup keys for the actual style-lookup table
names_aliases = {
    'b': 'blue',
    'g': 'green',
    'm': 'magenta',
    'o': 'orange',
    'p': 'purple',
    'r': 'red',
    'u': 'underline',

    'bb': 'bblue',
    'bg': 'bgreen',
    'bm': 'bmagenta',
    'bo': 'borange',
    'bp': 'bpurple',
    'br': 'bred',
    'bu': 'bunderline',

    'bb': 'bblue',
    'gb': 'bgreen',
    'mb': 'bmagenta',
    'ob': 'borange',
    'pb': 'bpurple',
    'rb': 'bred',
    'ub': 'bunderline',

    'hi': 'inverse',
    'inv': 'inverse',
    'mag': 'magenta',

    'flip': 'inverse',
    'swap': 'inverse',

    'reset': 'plain',
    'highlight': 'inverse',
    'hilite': 'inverse',
    'invert': 'inverse',
    'inverted': 'inverse',
    'swapped': 'inverse',

    'blueback': 'bblue',
    'grayback': 'bgray',
    'greenback': 'bgreen',
    'magback': 'bmagenta',
    'magentaback': 'bmagenta',
    'orangeback': 'borange',
    'purpleback': 'bpurple',
    'redback': 'bred',

    'bgblue': 'bblue',
    'bggray': 'bgray',
    'bggreen': 'bgreen',
    'bgmag': 'bmagenta',
    'bgmagenta': 'bmagenta',
    'bgorange': 'borange',
    'bgpurple': 'bpurple',
    'bgred': 'bred',

    'bluebg': 'bblue',
    'graybg': 'bgray',
    'greenbg': 'bgreen',
    'magbg': 'bmagenta',
    'magentabg': 'bmagenta',
    'orangebg': 'borange',
    'purplebg': 'bpurple',
    'redbg': 'bred',

    'backblue': 'bblue',
    'backgray': 'bgray',
    'backgreen': 'bgreen',
    'backmag': 'bmagenta',
    'backmagenta': 'bmagenta',
    'backorange': 'borange',
    'backpurple': 'bpurple',
    'backred': 'bred',
}

# names2styles matches color/style names to their ANSI-style strings
names2styles = {
    'blue': '\x1b[38;5;26m',
    'bold': '\x1b[1m',
    'gray': '\x1b[38;5;248m',
    'green': '\x1b[38;5;29m',
    'inverse': '\x1b[7m',
    'magenta': '\x1b[38;5;165m',
    'orange': '\x1b[38;5;166m',
    'plain': '\x1b[0m',
    'purple': '\x1b[38;5;99m',
    'red': '\x1b[31m',
    'underline': '\x1b[4m',

    'bblue': '\x1b[48;5;26m\x1b[38;5;15m',
    'bgray': '\x1b[48;5;248m\x1b[38;5;15m',
    'bgreen': '\x1b[48;5;29m\x1b[38;5;15m',
    'bmagenta': '\x1b[48;5;165m\x1b[38;5;15m',
    'borange': '\x1b[48;5;166m\x1b[38;5;15m',
    'bpurple': '\x1b[48;5;99m\x1b[38;5;15m',
    'bred': '\x1b[41m\x1b[38;5;15m',
}


def restyle_line(w, line: str, style: str) -> None:
    'Alternate styles for runs of digits in the string given.'

    start = 0
    end = len(line)
    if end > 1 and line[end - 2] == '\r' and line[end - 1] == '\n':
        end -= 2
    elif end > 0 and line[end - 1] == '\n':
        end -= 1

    while True:
        # see if line is over
        if start >= end:
            w.write('\n')
            return

        # find where the next run of digits starts; a negative index means
        # none were found
        i = -1
        for j in range(start, end):
            if line[j].isdigit():
                i = j
                break

        # check if rest of the line has no more digits
        if i < 0:
            w.write(line[start:end])
            w.write('\n')
            return

        # some ANSI-style sequences use 4-digit numbers, which are long
        # enough for this script to mangle
        is_ansi = i >= 2 and line[i-2] == '\x1b' and line[i-1] == '['

        # emit line up to right before the next run of digits starts
        w.write(line[start:i])
        start = i

        # find where/if the current run of digits ends; a negative index
        # means the run reaches the end of the line
        i = -1
        for j in range(start, end):
            if not line[j].isdigit():
                i = j
                break

        # check if rest of the line has only digits in it
        if i < 0:
            if not is_ansi:
                restyle_digits(w, line, start, end, style)
            else:
                w.write(line[start:end])
            w.write('\n')
            return

        # emit digits using alternate styling, and advance past them
        if not is_ansi:
            restyle_digits(w, line, start, i, style)
        else:
            w.write(line[start:i])
        start = i


def restyle_digits(w, digits: str, start: int, end: int, style: str) -> None:
    'Alternate styles on 3-item chunks from the string given.'

    diff = end - start

    # it's overall quicker to just emit short-enough digit-runs verbatim
    if diff < 4:
        w.write(digits[start:end])
        return

    # emit leading chunk of digits, which is the only one which
    # can have fewer than 3 items
    lead = diff % 3
    w.write(digits[start:start + lead])

    # the rest of the sub-string now has a multiple of 3 items left
    start += lead

    # start by styling the next digit-group only if there was a
    # non-empty leading group at the start of the full digit-run
    use_style = lead > 0

    # alternate styles until the string is over
    while start < end:
        # the digits left are always a multiple of 3
        stop = start + 3

        if use_style:
            w.write(style)
            w.write(digits[start:stop])
            w.write('\x1b[0m')
        else:
            w.write(digits[start:stop])

        # switch style and advance to the next 3-digit chunk
        use_style = not use_style
        start = stop


def seems_url(s: str) -> bool:
    protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
    return any(s.startswith(p) for p in protocols)


def handle_lines(w, src, style: str) -> None:
    for line in src:
        restyle_line(w, line, style)


args = argv[1:]
style = names2styles['gray']

# handle leading style/color option, if present
if len(args) > 0 and args[0].startswith('-'):
    s = args[0].lstrip('-')
    if s in names_aliases:
        s = names_aliases[s]

    if s in names2styles:
        style = names2styles[s]
        # skip leading arg, since it's clearly not a filepath
        args = args[1:]

if any(seems_url(e) for e in args):
    from urllib.request import urlopen

try:
    if args.count('-') > 1:
        msg = 'reading from `-` (standard input) more than once not allowed'
        raise ValueError(msg)

    for path in args:
        if path == '-':
            handle_lines(stdout, stdin, style)
            continue

        if seems_url(path):
            with urlopen(path) as inp:
                with TextIOWrapper(inp, encoding='utf-8') as txt:
                    handle_lines(stdout, txt, style)
            continue

        with open(path, encoding='utf-8') as inp:
            handle_lines(stdout, inp, style)

    if len(args) == 0:
        handle_lines(stdout, stdin, style)
except BrokenPipeError:
    # quit quietly, instead of showing a confusing error message
    stderr.close()
except KeyboardInterrupt:
    exit(2)
except Exception as e:
    print(f'\x1b[31m{e}\x1b[0m', file=stderr)
    exit(1)