import ast
import inspect


def get_tasklet_ast(stack_depth=2, frame=None) -> ast.With:
    """
    Returns the AST of the contents of the statement that called this function.
    :param stack_depth: If ``frame`` is None, how many levels up to go in the 
                        stack, default: 2 
                        (with statement->``__enter__``->``get_tasklet_ast``).
    :param frame: An optional argument specifying the frame, in order to avoid
                  inspecting it twice.
    :return: The AST of the body of the statement, or a FileNotFoundError if
             the AST cannot be recovered.
    """
    if frame is None:
        frame = inspect.stack()[stack_depth][0]
    caller = inspect.getframeinfo(frame)
    try:
        with open(caller.filename, 'r') as fp:
            pysrc = fp.read()
    except (OSError, FileNotFoundError):
        try:
            pysrc = inspect.getsource(frame)
        except (OSError, FileNotFoundError):
            raise FileNotFoundError('Cannot recover tasklet source code. This is likely because '
                                    'you are calling "with dace.tasklet:" from a Python terminal. '
                                    'Try to use Python code without tasklets instead, run from '
                                    'IPython, or a file.')

    module: ast.Module = ast.parse(pysrc)
    for node in ast.walk(module):
        if (getattr(node, 'lineno', -1) == caller.lineno and isinstance(node, ast.With)):
            return node

    raise FileNotFoundError('Cannot recover "with" statement from calling ' 'function.')