'''
Watiba pre-complier.  Watiba commands are BASH embedded commands between backtick characters (i.e. `), like traditional Bash captures.

Examples:
  example1.wt
    for line in `ls -lrt`.stdout:
        print(line)

    w = `cd /tmp && tar -zxvf blah.zip`
    if w.exit_code != 0:
        for l in w.stderr:
            print(l, file=stderr)

Author:
Ray Walker
raythonic@gmail.com

'''
import re
import sys

watiba_ref = "_watiba_"


# Singleton object.
class Compiler:
    def __init__(self, first_stmt):
        self.output = [first_stmt,
                       "import watiba",
                       f"{watiba_ref} = watiba.Watiba()"
                       ]
        self.resolver_count = 1
        self.spawn_call = []
        self.last_stmt = ""

        # Regex expressions for Watiba commands (order matters otherwise backticks would win over spawn)
        self.expressions = {
            # p = spawn `cmd`@host: block
            "^(\S.*)?spawn \s*`(\S.*)`@(\S.*) \s*?(\S.*)?:.*": self.spawn_generator_with_host,

            # p = spawn `cmd`: block
            "^(\S.*)?spawn \s*`(\S.*)`@(\S.*) \s*?(\S.*)?:.*": self.spawn_generator_with_host,

            # p = spawn `cmd`: block
            "^(\S.*)?spawn \s*`(\S.*)`\s*?(\S.*)?:.*": self.spawn_generator,

            # spawn-ctl {args}
            "^spawn-ctl \s*(\S.*)": self.spawn_ctl_args,

            # chain {host:cmd...
            "^(\S.*)?chain \s*`(\S.*)` \s*(\S.*)": self.chain_generator,

            # `cmd`@host
            ".*?([\-])?`(\S.*?)`@(\S.*) .*?": self.backticks_generator_with_host,

            # `cmd`
            ".*?([\-])?`(\S.*?)`.*?": self.backticks_generator
        }

    # Flush out any queue spawn calls that are located after the resolver block
    def flush(self):
        if len(self.spawn_call) > 0:
            if re.search("^return ", self.last_stmt.strip()):
                # Spit out spawn calls if they're queued up
                while len(self.spawn_call) > 0:
                    print(self.spawn_call.pop())
            else:
                print("ERROR in flush: Resolver block not properly terminated with return.", file=sys.stderr)
                print("    Block incorrectly terminated with:", file=sys.stderr)
                print(f"      {self.last_stmt}", file=sys.stderr)
                sys.exit(1)

    # Generate chain command
    def chain_generator(self, parms):
        assignment = parms["match"].group(1) if parms["match"].group(1) else ""
        cmd = f'"{parms["match"].group(2)}"' if parms["match"].group(2)[0] != "$" else parms["match"].group(2).replace("$", "")
        args = parms["match"].group(3)

        self.output.append(f'{assignment}_watiba_.chain({cmd}, {args})')

    # Set spawn controller args
    def spawn_ctl_args(self, parms):
        self.output.append(f'_watiba_.spawn_ctlr.set_parms({parms["match"].group(1)})')

    # Handle spawn code blocks (with host specified)
    def spawn_generator_with_host(self, parms):
        self.spawn_generator(parms, parms["match"].group(3))

    # Handle spawn code blocks
    def spawn_generator(self, parms, host=None):
        assign_idx = 1
        cmd_idx = 2
        args_idx = 3 if not host else 4

        # Build the spawn call that will be located just after the resolver block
        quote_style = "'" if "'" not in parms["match"].group(cmd_idx) else '"'

        # extract the command and if it's a variable, remove the $ and no quotes, otherwise in quotes
        cmd = parms["match"].group(cmd_idx)[1:] if parms["match"].group(cmd_idx)[
                                                       0] == "$" else f'{quote_style}{parms["match"].group(cmd_idx)}{quote_style}'
        # Build the next resolver method name
        resolver_name = f"__watiba_resolver_{self.resolver_count}__"
        self.resolver_count += 1

        # Include promise return if there's an assignment on the stmt
        promise_assign = parms["match"].group(assign_idx) if parms["match"].group(assign_idx) else ""

        # Add in args if there's any
        resolver_args = parms["match"].group(args_idx) if parms["match"].group(args_idx) else "{}"

        h = f'"{host}"' if host and host[0] != "$" else host
        h = h.replace("$", "") if h and h[0] == "$" else h

        # Queue up async call which is executed (spit out) at the end of the w_spawn block
        self.spawn_call.append(
            f'{parms["indentation"]}{promise_assign}_watiba_.spawn({cmd}, {resolver_name}, {resolver_args}, {"locals()"}, {h})')

        # Convert spawn `cmd`: statement to proper Python function definition
        self.output.append(f'{parms["indentation"]}def {resolver_name}(promise, args):')

    # Generator for `cmd` expressions
    def backticks_generator_with_host(self, parms):
        host = parms["match"].group(3)
        if host[0] == "$":
            host = host.replace("$", "")
        else:
            host = f'"{host}"'
        self.backticks_generator(parms, host)

    # Generator for `cmd` expressions
    def backticks_generator(self, parms, host=None):
        s = str(parms["statement"])

        # Run through the statement and replace ALL the backticked shell commands with Watiba function calls
        m = parms["match"]
        while m:
            # This flag control Watiba's CWD tracking
            context = False if m.group(1) == "-" else True
            cmd = m.group(2)

            # Make sure the string to be replaced includes the dash if needed
            repl_str = f'{"-" if not context else ""}`{cmd}`'

            # Replace the backticked commands with a Watiba function call
            if cmd[0] == "$":
                cmd = cmd[1:]
            else:
                quote_style = "'" if cmd.find("'") < 0 else '"'
                cmd = f"{quote_style}{cmd}{quote_style}"
            if not host:
                s = s.replace(repl_str, f"{watiba_ref}.bash({cmd}, {context})", 1)
            else:
                s = s.replace(repl_str, f"{watiba_ref}.ssh({cmd}, {host}, {context})", 1)
            # Test for more backticked commands
            m = re.search(parms["pattern"], s)

        self.output.append(s)

    # Compile the passed statement
    def compile(self, stmt):

        # Copy the statement to a local variable
        s = str(stmt)

        # Spit out spawn call if it's queued up (on block breaks)
        if len(s.strip()) > 0 and s.lstrip()[0] != "#" and len(s) - len(s.lstrip()) < len(self.last_stmt) - len(
                self.last_stmt.lstrip()):
            if len(self.spawn_call) > 0:
                if re.search("^return ", self.last_stmt.strip()):
                    print(self.spawn_call.pop())
                else:
                    print("ERROR: Resolver block not properly terminated with return.", file=sys.stderr)
                    print("    Block incorrectly terminated with:", file=sys.stderr)
                    print(f"      {self.last_stmt}", file=sys.stderr)
                    sys.exit(1)

        # Check the statement for a Watiba expresion
        for ex in self.expressions:
            m = re.search(ex, s.strip())

            # We have a Watiba expression. Generate the code.
            if m:
                return self.expressions[ex](
                    {"match": m,
                     "statement": s,
                     "pattern": ex,
                     "indentation": stmt[0:len(stmt) - len(stmt.lstrip())]
                     })

        self.output.append(stmt)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("ERROR. No input file.")
        sys.exit(0)

    # the versions array is generated at build time (see this module in bin/)
    if sys.argv[1] == "version":
        for v in versions:
            print(v)
        sys.exit(0)

    in_file = sys.argv[1]
    if not re.match(r".*\.wt$", in_file):
        print(f"ERROR: Input file must be type .wt.  Found {in_file}")
        sys.exit(1)

    out_file = in_file.replace('.wt', '.py', 1)
    c = None

    # Statements to ignore when looking for block terminations
    nothingness = ["#"]

    with open(in_file, 'r') as f:
        for statement in f:
            if not c:
                c = Compiler(statement.rstrip())
            else:
                c.compile(statement.rstrip())
                for o in c.output:
                    print(o)
                c.output = []
                if len(statement.strip()) > 0:
                    c.last_stmt = statement if len(statement.strip()) > 0 and statement.lstrip()[
                        0] not in nothingness else c.last_stmt

    # Flush out any queued spawn statement calls
    c.flush()
