"""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', 'check_types',
           'extract_commands']

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

# %% ../nbs/00_bashxtract.ipynb
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
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
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
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
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', 'Then', 'Else', 'Do', 'Loop'): visit(n.get(k))
        for k in ('Stmts', 'Items', 'Cases'): 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
_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
HANDLED_TYPES = {'File', 'CallExpr', 'BinaryCmd', 'Subshell', 'Block', 'IfClause', 'WhileClause', 
    'ForClause', 'CaseClause', 'CaseItem', 'FuncDecl', 'DeclClause', 'Lit', 'SglQuoted', 'DblQuoted',
    'ParamExp', 'CmdSubst', 'ProcSubst', 'Hdoc', 'Word', '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
def extract_commands(cmd, shfmt='shfmt'):
    "Split bash command into (commands, operators)"
    ast = parse_bash(cmd, shfmt=shfmt)
    check_types(ast)
    return visit_stmts(ast.get('Stmts', []), cmd), collect_ops(ast)
