#!/usr/bin/env python3
#
import argparse
import logging
import re
import sys
# noinspection PyUnresolvedReferences
from collections import defaultdict, Counter
from collections.abc import Iterable, Callable
from contextlib import contextmanager
# noinspection PyUnresolvedReferences
from datetime import *
from itertools import islice
# noinspection PyUnresolvedReferences
from math import *
from re import *

# logging
logging.basicConfig(level=logging.INFO, format='%(message)s', stream=sys.stderr)
logger = logging.getLogger(__name__)

# names that can be imported
available_names = {"Path": "pathlib",
                   "sleep": "time",
                   "randint": "random",
                   "get": "requests",
                   "b64encode": "base64",
                   "b64decode": "base64"}
# XX
#                    "reader": "csv",
#                    "writer": "csv"
# , "csv"
available_modules = set(list(available_names.values()) + ["jsonpickle", "humanize", "webbrowser", "collections"])

__doc__ = (f"Launch your tiny Python script on a piped in contents and pipe it out"
           "\n"
           "\nAvailable without import:"
           f"\n Loaded: re.* (match, search, findall), math.* (sqrt,...), datetime.* (datetime.now, ...),"
           f" defaultdict"
           f"\n Auto-imported functions: {', '.join(sorted(available_names.keys(), key=str.casefold))}"
           f"\n Auto-imported modules: {', '.join(sorted(available_modules))}"
           f"\n"
           f"\nAvailable variables:"
           f"\n * s – current line"
           f"\n * n – current line converted to an `int` (or `float`) if possible"
           f"\n * text – whole text, all lines together"
           f"\n * lines – list of lines so far processed"
           f"\n * numbers – list of numbers so far processed"
           f"\n * skip – omit line if True"
           f"\n * i=0, S=set(), L=list(), D=dict(), C=Counter() – other global variables"
           )

# parse arguments
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("command", help='Any Python script (multiple statements allowed)',
                    metavar="COMMAND", nargs="?")
parser.add_argument("-s", "--setup", help='Any Python script, executed before processing.'
                                          ' Useful for variable initializing.', metavar="COMMAND")
parser.add_argument("--finally", help='Any Python script, executed after processing.'
                                      ' Useful for final output.', metavar="COMMAND")
parser.add_argument("-v", "--verbose", help='Show command exceptions.'
                                            ' Used twice to show automatic imports.', action='count', default=0)
parser.add_argument("-f", "--filter", help='Line is piped out unchanged, however only if evaluated to True.',
                    action='store_true')
parser.add_argument("-n", help='Process only such number of lines.', type=int, metavar="NUM")
parser.add_argument("-1", help='Process just first line. Useful in combination with --whole.'
                               ' You may want to add -1 flag.', action='store_true')
parser.add_argument("-w", "--whole", help='Wait till whole text and then process.'
                                          ' Variable `text` is available containing whole text.', action='store_true')
parser.add_argument("--empty", help='Output empty lines. (By default skipped.)', action='store_true')
parser.add_argument("-0", help='Skip all lines output. (Useful in combination with --finally.)', action='store_true')
parser.add_argument("--lines", help='Populate `lines` and `numbers` with lines. This is off by default since this would'
                                    ' cause an overflow when handling an infinite input.', action='store_true')

regular = parser.add_argument_group("Regular output")
regular.add_argument("--search", help='Equivalent to `search(COMMAND, s)`', action='store_true')
regular.add_argument("--match", help='Equivalent to `match(COMMAND, s)`', action='store_true')
regular.add_argument("--findall", help='Equivalent to `findall(COMMAND, s)`', action='store_true')
regular.add_argument("--sub", help='Equivalent to `sub(COMMAND, SUBSTITUTION, s)`', metavar="SUBSTITUTION")

# evaluate command line arguments
args = parser.parse_args()
if getattr(args, "1"):
    args.n = 1

if not getattr(args, "command"):
    if getattr(args, "finally") is None:
        logger.error("You have to specify either COMMAND or --finally COMMAND.")
        quit()
    setattr(args, "0", True)  # suppress output if no command specified

if getattr(args, "finally"):
    args.lines = True

logger.setLevel({0: logging.ERROR, 1: logging.WARNING, 2: logging.INFO, 3: logging.DEBUG}[args.verbose])

# custom functions

match_class = match('', '').__class__ if sys.version_info < (3, 7) else Match  # drop with Python3.6


def write(s):
    """ Print either bytes or string. Bytes are not printed in the Python b-form: b'string' but raw. """
    if type(s) is bytes:
        sys.stdout.buffer.write(s + b'\n')
    else:
        print(s)


@contextmanager
def auto_import():
    """ If line processing failes with a NameError, check """
    global whole_hint_printed
    try:
        yield
    except NameError as e:
        name = re.match(r"name '(.*?)' is not defined", str(e))[1]
        if name:
            # Import anything on the fly (saved performance when loaded)
            if name == "text":
                if not whole_hint_printed:
                    logger.error("Did not you forget to use --whole to access `text`?")
                    whole_hint_printed = True
                raise
            elif name in ("numbers", "lines"):
                if not whole_hint_printed:
                    logger.error("Did not you forget to use --lines to access `lines` and `numbers`?")
                    whole_hint_printed = True
                raise
            elif name in available_names:
                module = available_names[name]
                logger.info(f"Importing {name} from {module}")
                # ex sleep = getattr(module "time", "sleep")
                globals()[name] = getattr(__import__(module), name)
            elif name in available_modules:
                logger.info(f"Importing {name}")
                globals()[name] = __import__(name)
            else:
                raise


def try_argument(callable_, argument, var, cmd="command"):
    """ Try to pass an argument to a callable. Returns False if TypeError happened. """
    t = f"attempt to use `{var}` as the callable parameter in {cmd}: {command[cmd]}({argument})"
    try:
        output(callable_(argument))
        command[cmd] += f"({var})"
    except TypeError as e:
        logger.debug(f"Failed {t} with: {e}")
        return False
    else:
        logger.debug(f"Successful {t}")
        return True


def output(line, final_round=None):
    """ output one or more lines """
    global tried_to_correct_callable
    if isinstance(line, match_class):
        # replace with the tuple of the groups or whole matched string (if no group matched)
        line = line.groups() or line.group(0)
    if line:  # empty string makes no output
        if isinstance(line, (str, bytes)):
            write(line)
        elif isinstance(line, list):  # list is output as multiple lines
            [output(el) for el in line]
        elif isinstance(line, Iterable):  # tuple or generator (but not a string) gets joined
            write(", ".join(str(el) for el in line))
        elif isinstance(line, Callable):  # tuple or generator (but not a string) gets joined
            try:
                output(line())
            except TypeError:
                if tried_to_correct_callable and not final_round:
                    # this it not the first line of command nor the `--finally` clause,
                    # we have already been there without success
                    raise
                tried_to_correct_callable = True
                # ex: `| pz webbrowser.open` -> `| pz webbrowser.open(s)`
                # ex: `sqrt() takes exactly one argument (0 given)`
                # ex: `open() missing required argument 'file' (pos 1)'` (build-in open)
                # ex: `open() missing 1 required positional argument: 'url'` (webbrowser.open)
                # Unfortunately, there is no certain way to determine the wanted type.
                # The wording of TypeError exceptions specifying the type vary.
                # The best we have is to use the inspect module to get the annotation or the parameter name.
                # We content to try it multiple things to pass as an argument.
                attempts = []
                if not final_round:
                    attempts.append((original_line, "s"))
                    if args.lines:
                        attempts.append((numbers, "numbers"))
                    else:
                        logger.debug("Since `--lines` flag is off, we will not try `numbers`.")
                    if n is not None:
                        # ex: echo  5 | pz sqrt | pz round
                        attempts.append((n, "n"))
                    import inspect
                    try:
                        param = list(inspect.signature(line).parameters.values())[0]
                        if param.name != "iterable":
                            # ex: pz b64encode += (s.encode('utf-8'))
                            logger.debug("Let's try `s.encode('utf-8')` automatically too.")
                            attempts.append((original_line.encode("utf-8"), "s.encode('utf-8')"))
                    except ValueError:  # ex: `set.add` raises no signature found
                        pass
                else:  # we are in the `--finally` clause, original_line is empty, we use `lines` or `numbers` instead
                    if len(numbers) == len(lines):
                        # ex: echo -e "1\n2\n3\n4" | pz --finally sum
                        attempts.append((numbers, "numbers", "finally"))
                    # ex: echo -e "1\n2\n3\n4" | pz  --finally "' - '.join" ->  1 - 2 - 3 - 4
                    attempts.append((lines, "lines", "finally"))
                if not any(try_argument(line, *x) for x in attempts):
                    raise
        else:  # ex: int, str
            write(line)
    else:
        if args.empty or (line == 0 and line is not False):
            write(line)


def get_number(s):
    num = None
    try:
        # we prefer having int over float because adding values '5' + '5' as '10'
        # looks better than '10.0' in most use cases
        num = float(s)
        num = int(s)  # "10.0" -> int conversion fails and num stays float
    except (ValueError, TypeError):
        pass
    return num


# variables available in the user scope
i = 0
S = set()
L = list()
D = dict()
C = Counter()
skip = None  # if user sets to False, the line will not be output
skip_all = getattr(args, "0")
whole_hint_printed = False

# prepare commands (prepend `line =` if needed)
command = {}
regular_command = None  # prepare regular modifications
reg_ex = None
if args.match or args.findall or args.search or args.sub:
    reg_ex = re.compile(args.command)
    if args.match:
        regular_command = reg_ex.match
    elif args.search:
        regular_command = reg_ex.search
    elif args.findall:
        regular_command = reg_ex.findall
    elif args.sub:
        regular_command = lambda line: reg_ex.sub(args.sub, line)


def prepare_command(name):
    command[name] = cmd = (getattr(args, name) or "").strip()

    if name == "command" and regular_command:
        # prepending `line = ` is not needed, the string is treated as a `match` parameter
        return

    # check if there is only a single line with a missing assignment
    if (len(cmd.splitlines()) == 1
            and not search("(s|skip)\s?[^=]?=[^=]", cmd)  # ex:
            # and not search("\(s\)", cmd)  # ex: L.append(s)
            and not any(cmd.lstrip().startswith(keyword) for keyword in ("if", "while", "for"))
            and ";" not in cmd
            and "lines." not in cmd):
        # "s = 1" - will not pass
        # "s += 1" - will not pass
        # "s + 1" - will pass
        # "s == 1" - will pass
        # X "(s)" - will not pass (ex: L.append(s))
        command[name] = ("skip = not " if args.filter else "s = ") + cmd
        logger.debug(f"Changing the {name} clause to: {command[name]}")


[prepare_command(name) for name in ("command", "finally")]

# prepare text processing (either fetch whole or line by line)
# Note: do not initialize the `text` to None. We want to be able to catch
# `<class 'NameError'> name 'text' is not defined` while not turning `--whole` on
text: str
if args.whole:
    # fetch whole text
    text = sys.stdin.read()  # XX we may stop reading text and use what we have on KeyboardInterrupt
    loop = (line for line in text.splitlines()[:args.n])
else:
    # load lines one by one (while taking at most N lines)
    loop = islice(sys.stdin, args.n)

# run setup
if args.setup:
    exec(args.setup)

# run processing
tried_to_correct_callable = False
original_line = None
n = None
lines: list
numbers: list
if args.lines:
    lines = []
    numbers = []

while True:
    try:
        original_line = s = next(loop).rstrip()
        n = get_number(s)
        if args.lines:
            lines.append(s)
            if n:
                numbers.append(n)
    except (KeyboardInterrupt, StopIteration):
        break

    try:
        while True:
            with auto_import():
                # loop until all on the fly imports are done
                skip = None
                # we process either a regular expression or a custom command
                if regular_command:
                    try:
                        s = regular_command(s)
                    except re.error as e:
                        logger.error(f"{e}, regular expression: {command['command']} on line: {s}")
                        break
                else:  # resolving custom command
                    # note that exec will not affect local field, hence we cannot easily put this in a method
                    exec(command["command"])
                if skip or (skip_all and skip is not False):  # user chooses to filter out the line
                    break
                output(s)
                break
    except Exception as e:
        logger.warning(f'Exception: {type(e)} {e} on line: {s}')
        continue

# run final script
if command["finally"]:
    original_line = s = n = None
    try:
        while True:
            with auto_import():
                exec(command["finally"])
                output(s, True)
                break
    except Exception as e:
        logger.warning(f'Exception: {type(e)} {e} in the --finally clause')
