"""
This module provides the generic functionality of tracing code by
modifying its AST. Eventually this will become a separate package.
This is similar to the standard library module bdb, while birdseye
itself would correspond to pdb.
Most of the work is in TreeTracerBase.
"""

from __future__ import print_function, division, absolute_import

from future import standard_library

standard_library.install_aliases()

import ast
import inspect
from collections import namedtuple
from copy import deepcopy
from functools import partial, update_wrapper
from itertools import takewhile
from typing import List, Dict, Any, Optional, NamedTuple, Tuple, Iterator, Callable, cast, Union
from types import FrameType, TracebackType, CodeType, FunctionType

from birdseye.utils import of_type, safe_next, PY3, Type, is_lambda, lru_cache, read_source_file, is_ipython_cell, \
    is_future_import


class TracedFile(object):
    """
    An instance of this class corresponds to a single .py file.
    It contains some useful data in the following attributes:

    - filename: name of the source file
    - source: textual contents of the file
    - root: root of the original Abstract Syntax Tree (AST) of the source,
            where the nodes of this tree have an additional handy attribute:
            - parent: parent of the node, so this node is a child node of its parent
    - tracer: instance of TreeTracerBase
    - code: executable code object compiled from the modified AST
    """

    is_ipython_cell = False

    def __init__(self, tracer, source, filename, flags):
        # type: (TreeTracerBase, str, str, int) -> None
        # Here the source code is parsed, modified, and compiled
        root = compile(source, filename, 'exec', ast.PyCF_ONLY_AST | flags, dont_inherit=True)  # type: ast.Module

        self.nodes = []  # type: List[ast.AST]

        def set_basic_node_attributes():
            self.nodes = []  # type: List[ast.AST]
            for node in ast.walk(root):  # type: ast.AST
                for child in ast.iter_child_nodes(node):
                    child.parent = node
                node._tree_index = len(self.nodes)
                self.nodes.append(node)

            # Mark __future__ imports and anything before (i.e. module docstrings)
            # to be ignored by the AST transformer
            for i, stmt in enumerate(root.body):
                if is_future_import(stmt):
                    for s in root.body[:i + 1]:
                        for node in ast.walk(s):
                            node._visit_ignore = True

        set_basic_node_attributes()

        new_root = tracer.parse_extra(root, source, filename)
        if new_root is not None:
            root = new_root

        set_basic_node_attributes()

        new_root = deepcopy(root)
        new_root = _NodeVisitor().visit(new_root)

        self.code = compile(new_root, filename, "exec", dont_inherit=True, flags=flags)  # type: CodeType
        self.root = root
        self.tracer = tracer
        self.source = source
        self.filename = filename


class FrameInfo(object):
    """
    Contains extra data about an execution frame.
    Can be obtained from the stack attribute of a TreeTracerBase instance
    """
    def __init__(self):
        # Stack of statements currently being executed
        self.statement_stack = []  # type: List[ast.stmt]

        # Stack of expression nodes within the above statement that
        # the interpreter is planning on evaluating, or has just evaluated
        # in the case of the last element of the list. For example, given
        # the expression f(g(x)), the stack would be [f, g, x] before and just
        # after evaluating x, since function arguments are evaluated before the
        # actual function call.
        self.expression_stack = []  # type: List[ast.expr]

        # Mapping from the expression node to its most recent value
        # in the corresponding frame
        self.expression_values = {}  # type: Dict[ast.expr, Any]

        # Node where the frame has explicitly returned
        # There may be parent nodes such as enclosing loops that still need to finish executing
        self.return_node = None  # type: Optional[ast.Return]

        # Mapping from expression nodes of list/set/dictionary comprehensions
        # which create a new frame of execution without a function call
        # to the most recent such frame.
        self.comprehension_frames = {}  # type: Dict[ast.expr, FrameType]

        # Most recent exception raised in the frame
        self.exc_value = None  # type: Optional[BaseException]


# Some of the attributes of the classes below are unused for now and are
# intended for future use, possibly by other debuggers


# Argument of TreeTracerBase.enter_call
EnterCallInfo = NamedTuple('EnterCallInfo', [

    # Node  from where the call was made
    ('call_node', Optional[Union[ast.expr, ast.stmt]]),

    # Node where the call begins
    ('enter_node', ast.AST),

    # Frame from which the call was made
    ('caller_frame', FrameType),

    # Frame of the call
    ('current_frame', FrameType)])

# Argument of TreeTracerBase.exit_call
ExitCallInfo = NamedTuple('ExitCallInfo', [

    # Node  from where the call was made
    ('call_node', Optional[Union[ast.expr, ast.stmt]]),

    # Node where the call explicitly returned
    ('return_node', Optional[ast.Return]),

    # Frame from which the call was made
    ('caller_frame', FrameType),

    # Frame of the call
    ('current_frame', FrameType),

    # Node where the call explicitly returned
    ('return_value', Any),

    # Exception raised in the call causing it to end,
    # will propagate to the caller
    ('exc_value', Optional[Exception]),

    # Traceback corresponding to exc_value
    ('exc_tb', Optional[TracebackType])])

# see TreeTracerBase.after_expr
ChangeValue = namedtuple('ChangeValue', 'value')


class TreeTracerBase(object):
    """
    Create a subclass of this class with one or more of the 'hooks'
    (methods which are empty in this class) overridden to take a custom action
    in the given situation. Then decorate functions with an instance of this class
    to trace them.
    """

    # Classes of expression nodes corresponding to expressions which create a new frame
    # of execution when evaluated, but are still evaluated immediately in the parent frame.
    SPECIAL_COMPREHENSION_TYPES = (ast.DictComp, ast.SetComp)  # type: Tuple[Type[ast.expr], ...]
    if PY3:
        SPECIAL_COMPREHENSION_TYPES += (ast.ListComp,)

    def __init__(self):
        # Mapping from frames of execution being traced to FrameInfo objects
        # for extra metadata.
        self.stack = {}  # type: Dict[FrameType, FrameInfo]

    @lru_cache()
    def compile(self, source, filename, flags=0):
        # type: (str, str, int) -> TracedFile
        return TracedFile(self, source, filename, flags)

    def _trace_methods_dict(self, traced_file):
        # type: (TracedFile) -> Dict[str, Callable]
        return {f.__name__: partial(f, traced_file)
                for f in [
                    self._treetrace_hidden_with_stmt,
                    self._treetrace_hidden_before_expr,
                    self._treetrace_hidden_after_expr,
                ]}

    def trace_function(self, func):
        # type: (FunctionType) -> FunctionType
        """
        Returns a version of the passed function with the AST modified to
        trigger the tracing hooks.
        """
        try:
            if inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func):
                raise ValueError('You cannot trace async functions')
        except AttributeError:
            pass

        if is_lambda(func):
            raise ValueError('You cannot trace lambdas')

        filename = inspect.getsourcefile(func)  # type: str

        if is_ipython_cell(filename):
            # noinspection PyPackageRequirements
            from IPython import get_ipython
            import linecache

            flags = get_ipython().compile.flags
            source = ''.join(linecache.cache[filename][2])
        else:
            source = read_source_file(filename)
            flags = 0

        # We compile the entire file instead of just the function source
        # because it can contain context which affects the function code,
        # e.g. enclosing functions and classes or __future__ imports
        traced_file = self.compile(source, filename, flags)

        # Then we have to recursively search through the newly compiled
        # code to find the code we actually want corresponding to this function
        code_options = []  # type: List[CodeType]

        def find_code(root_code):
            # type: (CodeType) -> None
            for const in root_code.co_consts:  # type: CodeType
                if not inspect.iscode(const):
                    continue
                matches = (const.co_firstlineno == func.__code__.co_firstlineno and
                           const.co_name == func.__code__.co_name)
                if matches:
                    code_options.append(const)
                find_code(const)

        find_code(traced_file.code)

        if len(code_options) > 1:
            # Currently lambdas aren't allowed anyway, but should be in the future
            assert is_lambda(func)
            raise ValueError("Failed to trace lambda. Convert the function to a def.")
        new_func_code = code_options[0]  # type: CodeType

        # Give the new function access to the hooks
        # We have to use the original __globals__ and not a copy
        # because it's the actual module namespace that may get updated by other code
        func.__globals__.update(self._trace_methods_dict(traced_file))

        # http://stackoverflow.com/a/13503277/2482744
        # noinspection PyArgumentList
        new_func = FunctionType(new_func_code, func.__globals__, func.__name__, func.__defaults__, func.__closure__)
        update_wrapper(new_func, func)  # type: FunctionType
        if PY3:
            new_func.__kwdefaults__ = getattr(func, '__kwdefaults__', None)
        new_func.traced_file = traced_file
        return new_func

    def __call__(self, func):
        # type: (FunctionType) -> FunctionType
        if inspect.isclass(func):
            raise TypeError('Decorating classes is no longer supported')

        return self.trace_function(func)

    def _treetrace_hidden_with_stmt(self, traced_file, _tree_index):
        # type: (TracedFile, int) -> _StmtContext
        """
        Called directly from the modified code.
        Every statement in the original code becomes:

        with _treetrace_hidden_with_stmt(...):
            <statement>
        """
        node = traced_file.nodes[_tree_index]
        node = cast(ast.stmt, node)
        frame = inspect.currentframe().f_back  # type: FrameType
        return _StmtContext(self, node, frame)

    def _treetrace_hidden_before_expr(self, traced_file, _tree_index):
        # type: (TracedFile, int) -> ast.expr
        """
        Called directly from the modified code before an expression is
        evaluated.
        """
        node = traced_file.nodes[_tree_index]
        node = cast(ast.expr, node)
        frame = inspect.currentframe().f_back  # type: FrameType

        frame_info = self.stack.get(frame)
        if frame_info is None:
            # This means there are no statements in this frame, and it's just started
            frame_info = FrameInfo()
            self.stack[frame] = frame_info

            # This is likely inside a comprehension. Find the parent frame corresponding
            # to a normal function call
            owner_frame = frame
            while owner_frame.f_code.co_name in ('<listcomp>', '<dictcomp>', '<setcomp>'):
                owner_frame = owner_frame.f_back
            if owner_frame != frame:
                comprehension = safe_next(of_type(self.SPECIAL_COMPREHENSION_TYPES,
                                                  ancestors(node)))  # type: ast.expr
                self.stack[owner_frame].comprehension_frames[comprehension] = frame

        frame_info.expression_stack.append(node)

        self.before_expr(node, frame)
        return node

    def _treetrace_hidden_after_expr(self, _, node, value):
        # type: (TracedFile, ast.expr, Any) -> Any
        """
        Called directly from the modified code after an expression is
        evaluated.
        """
        frame = inspect.currentframe().f_back  # type: FrameType
        result = self._after_expr(node, frame, value, None, None)
        if result is not None:
            assert isinstance(result, ChangeValue), "after_expr must return None or an instance of ChangeValue"
            value = result.value
        return value

    def _after_expr(self, node, frame, value, exc_value, exc_tb):
        frame_info = self.stack[frame]
        frame_info.expression_stack.pop()
        frame_info.expression_values[node] = value
        return self.after_expr(node, frame, value, exc_value, exc_tb)

    def _enter_call(self, enter_node, current_frame):
        # type: (ast.AST, FrameType) -> None
        caller_frame, call_node = self._get_caller_stuff(current_frame)
        self.stack[current_frame] = FrameInfo()
        self.enter_call(EnterCallInfo(call_node, enter_node, caller_frame, current_frame))

    def _get_caller_stuff(self, frame):
        # type: (FrameType) -> Tuple[FrameType, Optional[Union[ast.expr, ast.stmt]]]
        caller_frame = frame.f_back
        call_node = None
        if caller_frame in self.stack:
            frame_info = self.stack[caller_frame]
            expression_stack = frame_info.expression_stack
            if expression_stack:
                call_node = expression_stack[-1]
            else:
                call_node = frame_info.statement_stack[-1]  # type: ignore
        return caller_frame, call_node

    def _inner_node_and_frame(self, frame):
        """
        Given a normal frame with a nonempty expression stack, returns the actual
        frame and top of the expression stack after accounting for comprehensions
        that have their own frames.
        """
        frame_info = self.stack[frame]
        expression_stack = frame_info.expression_stack
        while isinstance(expression_stack[-1], TreeTracerBase.SPECIAL_COMPREHENSION_TYPES):
            frame = frame_info.comprehension_frames[expression_stack[-1]]
            expression_stack = self.stack[frame].expression_stack
        return expression_stack[-1], frame

    # The methods below are hooks meant to be overridden in subclasses to take custom actions

    def before_expr(self, node, frame):
        # type: (ast.expr, FrameType) -> None
        """
        Called right before the expression corresponding to `node` is evaluated
        within `frame`.
        """

    def after_expr(self, node, frame, value, exc_value, exc_tb):
        # type: (ast.expr, FrameType, Any, Optional[BaseException], Optional[TracebackType]) -> Optional[ChangeValue]
        """
        Called right after the expression corresponding to `node` is evaluated
        within `frame`. `value` is the value of the expression, if it succeeded.
        If the evaluation raised an exception, exc_value will be the exception object
        and exc_tb the traceback.

        Return `ChangeValue(x)` to change the value of the expression as
        seen by the rest of the program from `value` to `x`.
        """

    def before_stmt(self, node, frame):
        # type: (ast.stmt, FrameType) -> None
        """
        Called right before the statement corresponding to `node` is executed
        within `frame`.
        """

    def after_stmt(self, node, frame, exc_value, exc_traceback, exc_node):
        # type: (ast.stmt, FrameType, Optional[BaseException], Optional[TracebackType], Optional[ast.AST]) -> Optional[bool]
        """
        Called right after the statement corresponding to `node` is executed
        within `frame`.
        If the statement raised an exception, exc_value will be the exception object,
        exc_tb the traceback, and exc_node the node where the exception was raised
        (either this statement or an expression within).

        Returning True will suppress any exception raised (as with __exit__ in general).
        """

    def enter_call(self, enter_info):
        # type: (EnterCallInfo) -> None
        """
        Called before a function call begins executing. For typical `def` functions,
        this is called before the `before_stmt` for to the first statement in the function.
        """

    def exit_call(self, exit_info):
        # type: (ExitCallInfo) -> None
        """
        Called after a function call finishes executing. For typical `def` functions,
        this is called after the `after_stmt` for to the last statement to execute.
        """

    def parse_extra(self, root, source, filename):
        # type: (ast.Module, str, str) -> Optional[ast.Module]
        """
        Called before the AST (root) is modified to let subclasses make additional changes first.
        """


class _NodeVisitor(ast.NodeTransformer):
    """
    This does the AST modifications that call the hooks.
    """

    def generic_visit(self, node):
        # type: (ast.AST) -> ast.AST
        if not getattr(node, '_visit_ignore', False):
            if (isinstance(node, ast.expr) and
                    not (hasattr(node, "ctx") and not isinstance(node.ctx, ast.Load)) and
                    not isinstance(node, getattr(ast, 'Starred', ()))):
                return self.visit_expr(node)
            if isinstance(node, ast.stmt):
                return self.visit_stmt(node)
        return super(_NodeVisitor, self).generic_visit(node)

    def visit_expr(self, node):
        # type: (ast.expr) -> ast.Call
        """
        each expression e gets wrapped like this:
            _treetrace_hidden_after_expr(_treetrace_hidden_before_expr(_tree_index), e)

        where the _treetrace_* functions are the corresponding methods with the
        TreeTracerBase and traced_file arguments already filled in (see _trace_methods_dict)
        """

        before_marker = self._create_simple_marker_call(node, TreeTracerBase._treetrace_hidden_before_expr)
        ast.copy_location(before_marker, node)

        after_marker = ast.Call(
            func=ast.Name(id=TreeTracerBase._treetrace_hidden_after_expr.__name__,
                          ctx=ast.Load()),
            args=[
                before_marker,
                super(_NodeVisitor, self).generic_visit(node),
            ],
            keywords=[],
        )
        ast.copy_location(after_marker, node)
        ast.fix_missing_locations(after_marker)

        return after_marker

    def visit_stmt(self, node):
        # type: (ast.stmt) -> ast.With
        """
        Every statement in the original code becomes:

        with _treetrace_hidden_with_stmt(_tree_index):
            <statement>

        where the _treetrace_hidden_with_stmt function is the the corresponding method with the
        TreeTracerBase and traced_file arguments already filled in (see _trace_methods_dict)
        """
        context_expr = self._create_simple_marker_call(
            super(_NodeVisitor, self).generic_visit(node),
            TreeTracerBase._treetrace_hidden_with_stmt)

        if PY3:
            wrapped = ast.With(
                items=[ast.withitem(context_expr=context_expr)],
                body=[node],
            )
        else:
            wrapped = ast.With(
                context_expr=context_expr,
                body=[node],
            )
        ast.copy_location(wrapped, node)
        ast.fix_missing_locations(wrapped)
        return wrapped

    @staticmethod
    def _create_simple_marker_call(node, func):
        # type: (ast.AST, Callable) -> ast.Call
        """
        Returns a Call node representing `func(node._tree_index)`
        where node._tree_index is a numerical literal which allows the node object
        to be retrieved later through the nodes attribute of a TracedFile.
        """
        return ast.Call(
            func=ast.Name(id=func.__name__,
                          ctx=ast.Load()),
            args=[ast.Num(node._tree_index)],
            keywords=[],
        )


class _StmtContext(object):
    __slots__ = ('tracer', 'node', 'frame')

    def __init__(self, tracer, node, frame):
        # type: (TreeTracerBase, ast.stmt, FrameType) -> None
        self.tracer = tracer
        self.node = node
        self.frame = frame

    def __enter__(self):
        tracer = self.tracer
        node = self.node
        frame = self.frame
        if isinstance(node.parent, (ast.FunctionDef, ast.Module, ast.ClassDef)) and node is node.parent.body[0]:
            tracer._enter_call(node, frame)
        frame_info = tracer.stack[frame]
        frame_info.expression_stack = []
        frame_info.statement_stack.append(node)
        tracer.before_stmt(node, frame)

    def __exit__(self, exc_type, exc_val, exc_tb):
        # type: (Type[Exception], Exception, TracebackType) -> bool
        node = self.node
        tracer = self.tracer
        frame = self.frame
        frame_info = tracer.stack[frame]

        frame_info.statement_stack.pop()

        exc_node = None  # type: Optional[Union[ast.expr, ast.stmt]]
        if exc_val and exc_val is not frame_info.exc_value:
            exc_node = node
            frame_info.exc_value = exc_val

            # Call the after_expr hook if the exception was raised by an exception
            expression_stack = frame_info.expression_stack
            if expression_stack:
                exc_node, inner_frame = tracer._inner_node_and_frame(frame)
                tracer._after_expr(exc_node, inner_frame, None, exc_val, exc_tb)

        result = tracer.after_stmt(node, frame, exc_val, exc_tb, exc_node)

        if isinstance(node, ast.Return):
            frame_info.return_node = node

        parent = node.parent  # type: ast.AST
        return_node = frame_info.return_node
        exiting = (isinstance(parent, (ast.FunctionDef, ast.Module)) and
                   (node is parent.body[-1] or
                    exc_val or
                    return_node))
        if exiting:
            caller_frame, call_node = tracer._get_caller_stuff(frame)
            return_value = None
            if return_node and return_node.value and not exc_val:
                return_value = frame_info.expression_values[return_node.value]
            tracer.exit_call(ExitCallInfo(call_node,
                                          return_node,
                                          caller_frame,
                                          frame,
                                          return_value,
                                          exc_val,
                                          exc_tb
                                          ))

            del tracer.stack[frame]
            for comprehension_frame in frame_info.comprehension_frames.values():
                del tracer.stack[comprehension_frame]

        return result


def ancestors(node):
    # type: (ast.AST) -> Iterator[ast.AST]
    while True:
        try:
            node = node.parent
        except AttributeError:
            break
        yield node


Loop = Union[ast.For, ast.While, ast.comprehension]


def loops(node):
    # type: (ast.AST) -> Tuple[Loop, ...]
    """
    Return all the 'enclosing loops' of a node, up to the innermost class or
    function definition. This also includes the 'for in' clauses in list/dict/set/generator
    comprehensions. So for example, in this code:

      for x in ...:
          def foo():
              while True:
                  print([z for y in ...])

    The loops enclosing the node 'z' are 'while True' and 'for y in ...', in that order.
    """
    result = []
    while True:
        try:
            parent = node.parent
        except AttributeError:
            break
        if isinstance(parent, (ast.FunctionDef, ast.ClassDef)):
            break

        is_containing_loop = (((isinstance(parent, ast.For) and parent.iter is not node or
                                isinstance(parent, ast.While))
                               and node not in parent.orelse) or
                              (isinstance(parent, ast.comprehension) and node in parent.ifs))
        if is_containing_loop:
            result.append(parent)

        elif isinstance(parent, (ast.ListComp,
                                 ast.GeneratorExp,
                                 ast.DictComp,
                                 ast.SetComp)):
            generators = parent.generators
            if node in generators:
                generators = list(takewhile(lambda n: n != node, generators))
            result.extend(reversed(generators))

        node = parent

    result.reverse()
    return tuple(result)
