"""Extract commands used from bash command lines"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_bashxtract.ipynb.

# %% auto #0
__all__ = ['HANDLED_TYPES', 'parse_bash', 'part_text', 'word_text', 'nested_stmts', 'visit_stmts', 'collect_ops',
           'collect_redirects', 'scan_flag_args', 'check_types', 'extract_commands']

# %% ../nbs/00_bashxtract.ipynb #f67c4207
import shlex,subprocess,json,shutil
from fastcore.utils import *

# %% ../nbs/00_bashxtract.ipynb #3b43619d
def parse_bash(cmd:str, shfmt:str='shfmt'):
    "Parse `cmd` using `shfmt`"
    if not shutil.which(shfmt):
        raise FileNotFoundError(f"{shfmt} not found in PATH. If you installed safecmd with --user, add ~/.local/bin to your PATH.")
    res = subprocess.run([shfmt, '--to-json'], input=cmd, capture_output=True, text=True)
    if res.returncode != 0: raise ValueError(res.stderr)
    return json.loads(res.stdout)

# %% ../nbs/00_bashxtract.ipynb #9c322e10
def part_text(p, cmd):
    "Extracts the text value from a single word part node in the shfmt AST."
    t = p.get('Type', '')
    if t == 'Lit': return p['Value'].replace('\\ ', ' ')
    if t == 'SglQuoted': return p['Value']
    if t in ('DblQuoted', 'Hdoc'): return ''.join(part_text(x, cmd) for x in p.get('Parts', []))
    if t == 'ParamExp':
        name = p['Param']['Value']
        if p.get('Index'): name += '[' + word_text(p['Index'], cmd) + ']'
        return '${' + name + '}' if p.get('Rbrace') else '$' + name
    if t in ('CmdSubst', 'ProcSubst'):
        start, end = p['Pos']['Offset'], p['End']['Offset']
        return cmd[start:end]
    return ''

# %% ../nbs/00_bashxtract.ipynb #5379f8c2
def word_text(w, cmd):
    "Converts a Word node (with `Parts`) into its full text repr by concatenating `part_text` for each part."
    return ''.join(part_text(p, cmd) for p in w.get('Parts', []))

# %% ../nbs/00_bashxtract.ipynb #bef8b8c6
def nested_stmts(parts):
    "Yield all Stmts lists from nested Parts recursively"
    for p in parts:
        if stmts := p.get('Stmts'): yield stmts
        yield from nested_stmts(p.get('Parts', []))

# %% ../nbs/00_bashxtract.ipynb #4804fa0e
def visit_stmts(stmts, cmd, commands=None):
    "Visit statements, appending commands and handling redirects"
    if commands is None: commands=[]
    def visit(n):
        if not n: return
        if args := n.get('Args'):
            commands.append([word_text(a, cmd) for a in args])
            for a in args:
                for s in nested_stmts(a.get('Parts', [])): visit_stmts(s, cmd, commands)
        for k in ('Cmd', 'X', 'Y', 'Cond', 'Loop'): visit(n.get(k))
        for k in ('Stmts', 'Items', 'Cases', 'Then', 'Else', 'Do'): visit_stmts(n.get(k), cmd, commands)
    for s in stmts or []:
        visit(s)
        for r in s.get('Redirs', []):
            if not commands: continue
            if hdoc := r.get('Hdoc'): commands[-1].append(word_text(hdoc, cmd).rstrip('\n'))
            elif r.get('Op') == 63: commands[-1].extend(['<<<', word_text(r['Word'], cmd)])
    return commands

# %% ../nbs/00_bashxtract.ipynb #2ba597a5
_op_map = {10: '&&', 11: '||', 12: '|', 13: '|&', 54: '>', 55: '>>', 56: '<', 58: '<&', 59: '>&', 64: '&>', 65: '&>>'}
_attr_ops = {'Background': '&', 'Semicolon': ';', 'Assigns': '='}

def collect_ops(node, ops=None):
    "Walk AST node and collect all operators into a set"
    if ops is None: ops = set()
    if not isinstance(node, dict): return ops
    for attr, op in _attr_ops.items():
        if node.get(attr): ops.add(op)
    if op := _op_map.get(node.get('Op')): ops.add(op)
    for v in node.values(): collect_ops(v, ops) if isinstance(v, dict) else [collect_ops(x, ops) for x in v] if isinstance(v, list) else None
    return ops

# %% ../nbs/00_bashxtract.ipynb #3661batrbds
_write_ops = {54: '>', 55: '>>', 64: '&>', 65: '&>>'}

def collect_redirects(node, cmd, redirects=None):
    "Walk AST node and collect all write redirect destinations as (op, dest) tuples"
    if redirects is None: redirects = []
    if not isinstance(node, dict): return redirects
    for r in node.get('Redirs', []):
        if (op := _write_ops.get(r.get('Op'))) and (word := r.get('Word')):
            redirects.append((op, word_text(word, cmd)))
    for v in node.values():
        if isinstance(v, dict): collect_redirects(v, cmd, redirects)
        elif isinstance(v, list): [collect_redirects(x, cmd, redirects) for x in v]
    return redirects

# %% ../nbs/00_bashxtract.ipynb #c37daee0
def scan_flag_args(commands, exec_flags=None, dest_flags=None, dest_pos=None, exec_pos=None):
    "Scan commands for exec/dest flags and positional exec/dest args"
    extra_cmds, extra_dests = [], []
    if not exec_flags: exec_flags = {}
    if not dest_flags: dest_flags = {}
    if not dest_pos: dest_pos = {}
    if not exec_pos: exec_pos = {}
    
    for toks in commands:
        if not toks: continue
        cmd_name = toks[0]
        if cmd_name in exec_flags:
            for i, tok in enumerate(toks[:-1]):
                if tok in exec_flags[cmd_name]: extra_cmds.append(toks[i+1])
        if cmd_name in dest_flags:
            for i, tok in enumerate(toks[:-1]):
                if tok in dest_flags[cmd_name]: extra_dests.append((tok, toks[i+1]))
        args = toks[1:]
        if cmd_name in exec_pos:
            for idx in exec_pos[cmd_name]:
                try: extra_cmds.append(args[idx])
                except IndexError as e: print('exec pos not found', idx)
        if cmd_name in dest_pos:
            for idx in dest_pos[cmd_name]:
                try: extra_dests.append((idx, args[idx]))
                except IndexError as e: print('dest pos not found', idx)
    
    return extra_cmds, extra_dests


# %% ../nbs/00_bashxtract.ipynb #9e038c02
HANDLED_TYPES = {'File', 'CallExpr', 'BinaryCmd', 'Subshell', 'Block', 'IfClause', 'WhileClause', 
    'ForClause', 'CaseClause', 'CaseItem', 'FuncDecl', 'DeclClause', 'Lit', 'SglQuoted', 'DblQuoted',
    'ParamExp', 'CmdSubst', 'ProcSubst', 'Hdoc', 'Word', 'WordIter', 'Redirect', 'Comment', 'ArithmExp', 'ArithmCmd'}

def check_types(node):
    "Raise ValueError if AST contains unhandled node types"
    if not isinstance(node, dict): return
    if (t := node.get('Type')) and t not in HANDLED_TYPES: raise ValueError(f"Unhandled bash construct: {t}")
    for v in node.values():
        if isinstance(v, dict): check_types(v)
        elif isinstance(v, list): [check_types(x) for x in v]

# %% ../nbs/00_bashxtract.ipynb #9624cafa
def extract_commands(cmd, shfmt='shfmt', exec_flags=None, dest_flags=None, dest_pos=None, exec_pos=None):
    "Split bash command into (commands, operators, redirects)"
    ast = parse_bash(cmd, shfmt=shfmt)
    check_types(ast)
    commands = visit_stmts(ast.get('Stmts', []), cmd)
    ops = collect_ops(ast)
    redirects = collect_redirects(ast, cmd)
    
    extra_cmds, extra_dests = scan_flag_args(commands, exec_flags, dest_flags, dest_pos, exec_pos)
    for ecmd in extra_cmds:
        nested_cmds, nested_ops, nested_redirs = extract_commands(ecmd, shfmt, exec_flags, dest_flags, dest_pos, exec_pos)
        commands.extend(nested_cmds)
        ops.update(nested_ops)
        redirects.extend(nested_redirs)
    redirects.extend(extra_dests)
    
    return commands, ops, redirects

