#!/usr/bin/python

# The MIT License (MIT)
#
# Copyright (c) 2026 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.


info = '''
lpy [options...] [python expression] [files/URIs...]


Legend using PYthon colors lines by running the Python expression given on
each line, and using its common results. Inputs are encoded as UTF-8. Any
ANSI-styles are ignored from the input, to ensure consistent output-styling.

Carriage-returns are always ignored in lines, as well as any UTF-8-BOM on the
first line of each input.

The expression can use either `l` or `line` for the current line, and `i` as
a 0-based line counter which keeps growing even across input-sources, when
given multiple inputs. Also available is `n`, a 1-based line counter which
otherwise works the same way.

Each line is automatically parsed as JSON: when successful, the parsed line
is available to the expression as `v`, or `value`. You can check failure to
parse JSON by checking if `v` or `value` is of type Skip, since None can be
the result of successfully parsing a null JSON value.

Input-sources can be either files or web-URIs. When not given any explicit
named sources, the standard input is used. A single dash also means the
standard input.

When the expression results in None, the current input line is ignored,
allowing expressions to also act as convenient line filters.


Options

All options can start with either a single or a double leading dash:

    -h, -help                      show this help message
    -m, -mod, -module, -modules    import modules named in the next argument,
                                   where multiple names are comma-separated
    -t, -trace, -traceback         turn exceptions into multi-line tracebacks
    -v, -verbose                   also show legend-keys to stderr at the end


Examples

# read from a file, and only emit (equally-colored) lines exceeding 80 items
lpy '1 if len(line) > 80 else None' data.txt

# color files in current folder (roughly) by common auto-detected file-type
file * | lpy '(lambda e: e[1] + e[2])(" ".join(l.split(":")[1:]).split())'

# color files in current folder (roughly) by common auto-detected file-type
file * | lpy 'fields[1] if len(fields) < 3 else fields[1] + fields[2]'

# color numbers by common remainder when divided by 5
seq 50 | lpy 'v % 5'

# transform lines in order with a format-string, styling nothing
seq 50 | lpy 'print(f"{line + "**2":>6} = {float(v**2):12,.4f}")'

# color numbers by common remainder when divided by 5, ignoring numbers > 30
seq 50 | lpy 'None if v > 30 else v % 5'
'''


from itertools import islice, zip_longest
from json import dump, dumps, loads
from math import isinf, isnan
from os import getenv
from re import compile as compile_uncached, IGNORECASE
from sys import argv, exit, stderr, stdin, stdout
from time import localtime, strftime
from typing import Iterable


if len(argv) > 1 and argv[1] in ('-h', '--h', '-help', '--help'):
    print(info.strip())
    exit(0)


def style_lines(src, expr):
    # `comprehension` expressions seem to ignore local variables: even
    # lambda-based workaround-attempts fail to make needed values like
    # the current line available to such expressions
    global i, n, fn, l, line, v, val, value
    global legend, styles, resets, items, fields, words, tsv

    fn = 1

    for l in src:
        l = l.rstrip('\r\n').rstrip('\n')
        if fn == 1:
            l = l.lstrip('\xef\xbb\xbf')

        line = l
        items = fields = words = line.split()
        tsv = line.split('\t')

        try:
            v = val = value = loads(l)
        except Exception as _:
            v = val = value = Skip()

        k = eval(expr)
        if callable(k) and (not isinstance(k, Skip)):
            k = k(line)
        if isinstance(k, BaseException):
            raise k

        i += 1
        n += 1
        fn += 1

        if (k is None) or isinstance(k, Skip):
            continue
        if not (k in legend):
            legend[k] = len(legend)

        k = legend[k] % len(styles)
        # out_line = line.replace('\x1b[0m', resets[k])
        out_line = plain(line)
        print(f'{styles[k]}{out_line}\x1b[0m', flush=flushed)

def fail(msg, code = 1):
    print(str(msg), file=stderr)
    exit(code)

def make_open_utf8(open):
    def open_utf8_readonly(path):
        return open(path, encoding='utf-8')
    return open_utf8_readonly

def seemsurl(path):
    protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
    return any(path.startswith(p) for p in protocols)

class Skip:
    def __call__(self, x):
        return isinstance(x, self.__class__)

skip = Skip()

def cases(x, *args):
    '''
    Simulate a switch statement on a value, using matches/result pairs from
    the arguments given; when given an even number of extra args, None is
    used as a final fallback result; when given an odd number of extra args,
    the last argument is used as a final `default` value, if needed.
    '''

    for i in range(0, len(args) - len(args) % 2, 2):
        test, res = args[i], args[i+1]
        if isinstance(test, (list, tuple)) and (x in test):
            return res
        if isinstance(test, float) and isnan(test) and isnan(x):
            return res
        if x == test:
            return res
    return None if len(args) % 2 == 0 else args[-1]

switch = cases

def chunk(items, chunk_size):
    'Break iterable into chunks, each with up to the item-count given.'

    if isinstance(items, str):
        n = len(items)
        while n >= chunk_size:
            yield items[:chunk_size]
            items = items[chunk_size:]
            n -= chunk_size
        if n > 0:
            yield items
        return

    if not isinstance(chunk_size, int):
        raise Exception('non-integer chunk-size')
    if chunk_size < 1:
        raise Exception('non-positive chunk-size')

    it = iter(items)
    while True:
        head = tuple(islice(it, chunk_size))
        if not head:
            return
        yield head

chunked = chunk

# re_cache is used by custom func compile to cache previously-compiled
# regular-expressions, which makes them quicker to (re)use in formulas
re_cache = {}

def re_compile(expr, flags = 0):
    'Speed-up using regexes across lines, by avoiding recompilations.'

    if flags in re_cache:
        cache = re_cache[flags]
    else:
        cache = {}
        re_cache[flags] = cache
    if expr in cache:
        return cache[expr]

    pat = compile_uncached(expr, flags)
    cache[expr] = pat
    return pat

def icompile(expr):
    return re_compile(expr, IGNORECASE)

def cond(*args):
    for i, e in enumerate(args):
        if i % 2 == 0 and i < len(args) - 1 and e:
            return args[i + 1]

    return args[-1] if len(args) % 2 == 1 else None

def denan(x, fallback = None):
    'Replace floating-point NaN with the alternative value given.'
    return x if not (isinstance(x, float) and isnan(x)) else fallback

def denil(*args):
    'Avoid None values, if possible: first value which isn\'t None wins.'
    for e in args:
        if e != None:
            return e
    return None

denone = denull = denil

def dive(into, using):
    'Depth-first recursive caller for 1-input functions.'

    if callable(into):
        into, using = using, into

    def rec(v):
        if isinstance(v, dict):
            return {k: rec(v) for k, v in v.items()}
        if isinstance(v, Iterable) and not isinstance(v, str):
            return [rec(v) for v in v]
        return using(v)

    return rec(into)

def divekeys(into, using):
    'Depth-first recursive caller for 2-input funcs which rename dict keys.'

    if callable(into):
        into, using = using, into

    def rec(v):
        if isinstance(v, dict):
            return {using(k): rec(v) for k, v in v.items()}
        if isinstance(v, Iterable) and not isinstance(v, str):
            return [rec(v) for i, v in enumerate(v)]
        return v

    return rec(None, into)

def divekv(into, using, using2 = None):
    'Depth-first recursive caller for 2-input functions.'

    if using2 is None:
        if callable(into):
            into, using = using, into
    else:
        if not callable(using2):
            into, using, using2 = using2, into, using

    def rec(k, v):
        if isinstance(v, dict):
            return {k: rec(k, v) for k, v in v.items()}
        if isinstance(v, Iterable) and not isinstance(v, str):
            return [rec(i, v) for i, v in enumerate(v)]
        return using(k, v)

    def rec2(k, v):
        if isinstance(v, dict):
            return {str(using(k, v)): rec2(k, v) for k, v in v.items()}
        if isinstance(v, Iterable) and not isinstance(v, str):
            return [rec2(i, v) for i, v in enumerate(v)]
        return using2(k, v)

    return rec(None, into) if using2 is None else rec2(None, into)

kvdive = divekv

def drop(src, *what):
    if isinstance(src, str):
        for s in what:
            src = src.replace(s, '')
        return src

    def kdrop(src, what):
        return {k: v for (k, v) in src.items() if not (k in what)}

    if isinstance(src, dict):
        return kdrop(src, set(what))

    if isinstance(src, Iterable):
        what = set(what)
        return [kdrop(e, what) for e in src if isinstance(e, dict)]

    return None

dropped = drop

def flat(*args):
    'Turn arbitrarily-nested values/sequences into a single flat sequence.'

    def rec(x):
        if x is None:
            return

        if isinstance(x, dict):
            yield from rec(x.values())

        if isinstance(x, str):
            yield x
            return

        if isinstance(x, Iterable):
            for e in x:
                yield from rec(e)
            return

        yield x

    for x in args:
        yield from rec(x)

flatten = flattened = flat

def identity(x):
    return x

def iota(n):
    return range(1, n + 1)

def join(x, *y):
    'Join values into a string, or make a dict from keys and values.'

    if len(y) == 0:
        return ' '.join(str(v) for v in x)
    if isinstance(x, str):
        return x.join(str(v) for v in y)
    if len(y) == 1 and isinstance(y[0], str):
        return y[0].join(str(v) for v in x)
    if len(y) == 1 and isinstance(y[0], (list, range, set, tuple, Generator)):
        return {k: v for k, v in zip_longest(x, y[0]) if not (k is None)}
    if isinstance(y, (list, range, set, tuple, Generator)):
        return {k: v for k, v in zip_longest(x, y) if not (k is None)}
    return {k: y for k in x}

joined = join

def maybe(f, x = None):
    try:
        return f(x)
    except Exception as _:
        return x

def number(x):
    try:
        return int(x)
    except Exception as _:
        pass
    try:
        return float(x)
    except Exception as _:
        return x

def pick(src, *keys):
    if isinstance(src, dict):
        return {k: src.get(k, None) for k in keys}
    return [{k: e.get(k, None) for k in keys} for e in src if isinstance(e, dict)]

picked = pick

def plain(s):
    'Ignore all ANSI-style sequences in a string.'
    return re_compile('''\x1b\\[([0-9;]+m|[0-9]*[A-HJKST])''').sub('', s)

def predicate(x):
    'Helps various higher-order funcs, by standardizing `predicate` values.'
    if callable(x):
        return x
    if not isinstance(x, float):
        return lambda y: x == y
    if isnan(x):
        return lambda y: isinstance(y, float) and isnan(y)
    if isinf(x):
        return lambda y: isinstance(y, float) and isinf(y)

def rescue(attempt, fallback = None):
    try:
        return attempt()
    except Exception as e:
        if callable(fallback) and (not isinstance(fallback, Skip)):
            return fallback(e)
        return fallback

rescued = rescue

def retype(x):
    'Try to narrow the type of the value given.'

    if isinstance(x, float):
        n = int(x)
        return n if float(n) == x else x

    if not isinstance(x, str):
        return x

    try:
        return loads(x)
    except Exception as _:
        pass

    try:
        return int(x)
    except Exception as _:
        pass

    try:
        return float(x)
    except Exception as _:
        pass

    return x

autocast = autocasted = mold = molded = recast = recasted = remold = retype
remolded = retyped = retype

def json0(x):
    if isinstance(x, (range, set, Generator)):
        x = tuple(x)
    return dumps(x, separators=(',', ':'), allow_nan=False, indent=None)

j0 = json0

def typeof(x):
    return {
        type(None): 'null',
        bool: 'boolean',
        dict: 'object',
        float: 'number', int: 'number',
        str: 'string',
        list: 'array', tuple: 'array',
    }.get(type(x), 'other')

jstype = typeof

def after(x, what):
    i = x.find(what)
    return '' if i < 0 else x[i+len(what):]

def afterlast(x, what):
    i = x.rfind(what)
    return '' if i < 0 else x[i+len(what):]

afterfinal = afterlast

def before(x, what):
    i = x.find(what)
    return x if i < 0 else x[:i]

def beforelast(x, what):
    i = x.rfind(what)
    return x if i < 0 else x[:i]

beforefinal = beforelast

def since(x, what):
    i = x.find(what)
    return '' if i < 0 else x[i:]

def sincelast(x, what):
    i = x.rfind(what)
    return '' if i < 0 else x[i:]

sincefinal = sincelast

def until(x, what):
    i = x.find(what)
    return x if i < 0 else x[:i+len(what)]

def untilfinal(x, what):
    i = x.rfind(what)
    return x if i < 0 else x[:i+len(what)]

untillast = untilfinal

def message(msg, result = None):
    print(msg, file=stderr)
    return result

msg = message

# seen is used by func `once` to remember previously-given values
seen = set()

def once(x, key = None):
    k = key(x) if key else x
    if k in seen:
        return None
    seen.add(k)
    return x

dedup = unique = once

def utf8(x):
    try:
        if isinstance(x, str):
            x = x.encode('utf-8')
        return str(x, 'utf-8')
    except Exception as _:
        return None

def ymdhms(when = None):
    fmt = f'%Y-%m-%d %H:%M:%S'
    if isinstance(when, (float, int)):
        return strftime(fmt, localtime(float(when)))
    if isinstance(when, tuple):
        return strftime(fmt, when)
    return strftime(fmt, localtime())


amp = '&'
apos = '\''
lcur = lcurly = '{'
quot = '"'
rcur = rcurly = '}'

true = True
false = False
nil = none = null = None


exec = None
open_utf8 = make_open_utf8(open)
open = open_utf8

modules_opts = (
    '-m', '--m', '-mod', '--mod', '-module', '--module',
    '-modules', '--modules',
)
trace_opts = ('-t', '--t', '-trace', '--trace', '-traceback', '--traceback')
verbose_opts = ('-v', '--v', '-verbose', '--verbose')

args = argv[1:]
if any(seemsurl(e) for e in args):
    from io import TextIOWrapper
    from urllib.request import urlopen

tracing = False
verbose = False

while len(args) > 0:
    if args[0] == '-':
        break

    if args[0] == '--':
        args = args[1:]
        break

    if args[0] in modules_opts:
        try:
            if len(args) < 2:
                msg = 'a module name or a comma-separated list of modules'
                raise Exception('expected ' + msg)

            g = globals()
            from importlib import import_module
            for e in args[1].split(','):
                g[e] = import_module(e)

            g = None
            import_module = None
            args = args[2:]
        except Exception as e:
            fail(e, 1)

        continue

    if args[0] in verbose_opts:
        verbose = True
        args = args[1:]
        continue

    if args[0] in trace_opts:
        tracing = True
        args = args[1:]
        continue

    if args[0].startswith('-'):
        fail(f'lpy: unsupported option {args[0]}', 1)
        exit(1)

    break


# ensure live-lines output, unless stdout is being saved into a file
flushed = stdout.isatty() or (not stdout.seekable())

# '\x1b[38;2;255;135;255m', # pink (originally after cyan)

styles = (
    '\x1b[38;2;0;95;215m', # blue
    '\x1b[38;2;215;95;0m', # orange
    '\x1b[38;2;135;95;255m', # purple
    '\x1b[38;2;0;175;215m', # cyan
    '\x1b[38;2;255;85;255m', # pink
    '\x1b[38;2;0;135;95m', # green
    '\x1b[38;2;204;0;0m', # red
    '\x1b[38;2;168;168;168m', # gray
)

if (getenv('COLORBLIND', '0') != '0') or (getenv('COLOR_BLIND', '0') != '0'):
    styles = (
        '\x1b[38;2;0;95;215m', # blue
        '\x1b[38;2;215;95;0m', # orange
        '\x1b[38;2;135;95;255m', # purple
        '\x1b[38;2;0;175;215m', # cyan
        '\x1b[38;2;255;85;255m', # pink
        '\x1b[38;2;0;135;95m', # green
        '\x1b[38;2;168;168;168m', # gray
    )
getenv = None

resets = tuple('\x1b[0m' + s for s in styles)

expr = '.'
if len(args) > 0:
    expr = args[0]
    args = args[1:]

if expr == '.':
    expr = 'line'

try:
    expr = compile(expr, expr, mode='eval')
    compile = None

    i = 0
    n = 1
    v = val = value = Skip()

    legend = {}
    items = fields = words = []
    tsv = []

    if len(args) == 0:
        style_lines(stdin, expr)

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

    for path in args:
        if path == '-':
            style_lines(stdin, expr)
            continue

        if seemsurl(path):
            with urlopen(path) as inp:
                with TextIOWrapper(inp, encoding='utf-8') as txt:
                    style_lines(txt, expr)
            continue

        with open_utf8(path) as txt:
            style_lines(txt, expr)

    if verbose:
        for k, v in legend.items():
            print(f'{styles[v]}{k}\x1b[0m', file=stderr)
except BrokenPipeError:
    # quit quietly, instead of showing a confusing error message
    stderr.close()
    exit(0)
except KeyboardInterrupt:
    exit(2)
except Exception as e:
    if tracing:
        raise e
    fail(e, 1)