#!/usr/bin/python3

# The MIT License (MIT)
#
# Copyright © 2020-2025 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.


from itertools import islice
from math import isinf, isnan, modf, nan
from sys import argv, exit, stderr, stdin, stdout
from typing import Dict, Tuple


info = '''
ntsv [filepaths/URIs...]


Nice Tab Separated Values realigns and styles data tables using ANSI color
sequences. When not given filepaths/URIs to read data from, this tool reads
from standard input.
'''

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


# spaces is used by func emit_spaces
spaces = tuple(' ' * i for i in range(64))


def emit_spaces(w, n: int) -> None:
    'Speed-up writing spaces, by minimizing calls to func write.'

    if n < 1:
        return

    while n >= len(spaces):
        w.write(spaces[-1])
        n -= len(spaces)
    w.write(spaces[n])


def count_decimals(s: str) -> int:
    '''
    Count the trailing decimal digits in a string, excluding the decimal dot.
    This func assumes the string can parse as a valid decimal-format float.
    '''
    k = s.find('.')
    return len(s) - k - 1 if k >= 0 else 0


def get_table_info(src) -> Dict:
    'Gather metadata by reading TSV rows from the input given.'

    rows = []
    totals = []
    numeric = []
    max_decs = []
    max_widths = []

    def handle_row(items: Tuple[str]) -> None:
        nonlocal totals, numeric, max_decs, max_widths

        if len(max_widths) < len(items):
            extras = len(items) - len(max_widths)
            totals.extend([0.0] * extras)
            numeric.extend([False] * extras)
            max_decs.extend([0] * extras)
            max_widths.extend([0] * extras)

        for i, s in enumerate(items):
            max_widths[i] = max(max_widths[i], len(s))
            try:
                f = float(s)
            except Exception:
                continue

            if not (isnan(f) or isinf(f)):
                max_decs[i] = max(max_decs[i], count_decimals(s))
                totals[i] += f
                numeric[i] = True

    for line in src:
        line = line.rstrip('\r\n').rstrip('\n')
        if line == '':
            continue

        items = tuple(line.split('\t'))
        rows.append(items)
        handle_row(items)

    # format column-totals as strings, and update column widths with those
    for i, num in enumerate(numeric):
        s = f'{totals[i]:.{max_decs[i]}f}' if num else '-'
        # s = f'{totals[i]:.{max_decs[i]}f}' if num else ''
        max_widths[i] = max(max_widths[i], len(s))
        totals[i] = s

    return {
        'rows': rows,
        'max-widths': tuple(max_widths),
        'max-decimals': tuple(max_decs),
        'totals': tuple(totals),
    }


def show_table(w, meta: Dict) -> None:
    'Show/render table using the metadata given.'

    gap = len('  ')
    rows = meta['rows']
    max_widths = meta['max-widths']
    max_decs = meta['max-decimals']
    totals = meta['totals']

    if len(rows) == 0:
        return

    for row in rows:
        show_tiles(w, row, len(max_widths))
        show_row(w, row, max_widths, max_decs, gap)

    # make the totals line stand out by padding it with spaces, instead of
    # showing tiles: their lack should give a hint this final `row` isn't
    # part of the input
    emit_spaces(w, len(totals))
    show_row(w, totals, max_widths, max_decs, gap)


def show_row(w, row: Tuple[str], widths: Tuple[int], max_decs: Tuple[int],
                gap: int) -> None:
    pad = 0

    for i, s in enumerate(row):
        try:
            f = float(s)
        except Exception:
            f = nan

        if not (isnan(f) or isinf(f)):
            decs = count_decimals(s)
            if decs > 0:
                trail = max_decs[i] - decs
            elif max_decs[i] > 0:
                trail = max_decs[i] + 1
            else:
                trail = 0

            lead = widths[i] - len(s)
            emit_spaces(w, max(pad + lead - trail, 0) + gap)
            rem, _ = modf(f)
            if rem != 0:
                if f > 0:
                    w.write(f'\x1b[38;2;0;135;95m{s}\x1b[0m')
                elif f < 0:
                    w.write(f'\x1b[38;2;215;95;95m{s}\x1b[0m')
                elif f == 0:
                    w.write(f'\x1b[38;2;0;95;215m{s}\x1b[0m')
                else:
                    w.write(s)
            else:
                if f > 0:
                    w.write(f'\x1b[38;2;0;95;0m{s}\x1b[0m')
                elif f < 0:
                    w.write(f'\x1b[38;2;204;0;0m{s}\x1b[0m')
                elif f == 0:
                    w.write(f'\x1b[38;2;0;95;215m{s}\x1b[0m')
                else:
                    w.write(s)
            pad = trail
        else:
            emit_spaces(w, pad + gap)
            w.write(s)
            pad = widths[i] - len(s)

    # don't forget to end the row/line
    w.write('\n')


def show_tiles(w, row: Tuple[str], expected: int) -> None:
    for s in row:
        try:
            f = float(s)
            rem, _ = modf(f)
            if f > 0:
                w.write('\x1b[38;2;0;135;95m■' if rem != 0 else '\x1b[38;2;0;95;0m■')
            elif f < 0:
                w.write('\x1b[38;2;215;95;95m■' if rem != 0 else '\x1b[38;2;204;0;0m■')
            elif f == 0:
                w.write('\x1b[38;2;0;95;215m■')
            else:
                w.write('\x1b[0m■')
        except Exception:
            if s == '':
                w.write('\x1b[0m○')
            else:
                w.write('\x1b[38;2;128;128;128m■')

    if len(row) - expected:
        w.write('\x1b[0m')
    for _ in range(expected - len(row)):
        w.write('×')
    w.write('\x1b[0m')


def handle_tsv(w, src) -> None:
    show_table(w, get_table_info(src))


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


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

    if any(seems_url(e) for e in argv):
        from io import TextIOWrapper
        from urllib.request import urlopen

    for path in islice(argv, 1, None):
        if path == '-':
            handle_tsv(stdout, stdin)
        elif seems_url(path):
            with urlopen(path) as inp:
                with TextIOWrapper(inp, encoding='utf-8') as txt:
                    handle_tsv(stdout, txt)
        else:
            with open(path, encoding='utf-8') as inp:
                handle_tsv(stdout, inp)

    if len(argv) < 2:
        handle_tsv(stdout, stdin)
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)