"""Simple traceback introspection. Used to add additional information to
AssertionErrors in tests, so that failure messages may be more informative."""
import inspect
import logging
import re
import textwrap
import tokenize
import collections
from io import StringIO

log = logging.getLogger(__name__)


def inspect_traceback(tb):
    """Inspect a traceback and its frame, returning source for the expression
    where the exception was raised, with simple variable replacement performed
    and the line on which the exception was raised marked with '>>' """
    log.debug('inspect traceback %s', tb)
    while tb.tb_next:
        tb = tb.tb_next
    frame = tb.tb_frame
    lines, exc_line = tbsource(tb)
    inspect_lines, mark_line = find_inspectable_lines(lines, exc_line)
    src = StringIO(textwrap.dedent(''.join(inspect_lines)))
    exp = Expander(frame.f_locals, frame.f_globals)
    while inspect_lines:
        try:
            for tok in tokenize.generate_tokens(src.readline):
                exp(*tok)
        except tokenize.TokenError as e:
            log.debug("Tokenizer error: %s", e)
            inspect_lines.pop(0)
            mark_line -= 1
            src = StringIO(textwrap.dedent(''.join(inspect_lines)))
            exp = Expander(frame.f_locals, frame.f_globals)
            continue
        break
    padded = []
    if exp.expanded_source:
        exp_lines = exp.expanded_source.split('\n')
        ep = 0
        for line in exp_lines:
            if ep == mark_line:
                padded.append('>>  ' + line)
            else:
                padded.append('    ' + line)
            ep += 1
    return '\n'.join(padded)


def tbsource(tb, context=6):
    """Get source from  a traceback object.
    A tuple of two things is returned: a list of lines of context from
    the source code, and the index of the current line within that list.
    The optional second argument specifies the number of lines of context
    to return, which are centered around the current line."""
    lineno = tb.tb_lineno
    frame = tb.tb_frame
    if context > 0:
        start = lineno - 1 - context//2
        log.debug("lineno: %s start: %s", lineno, start)
        try:
            lines, dummy = inspect.findsource(frame)
        except IOError:
            lines, index = [''], 0
        else:
            all_lines = lines
            start = max(start, 1)
            start = max(0, min(start, len(lines) - context))
            lines = lines[start:start+context]
            index = lineno - 1 - start
            if index > 0:
                while lines[index-1].strip().endswith('\\'):
                    start -= 1
                    lines = all_lines[start:start+context]
    else:
        lines, index = [''], 0
    log.debug("tbsource lines '''%s''' around index %s", lines, index)
    return (lines, index)


def find_inspectable_lines(lines, pos):
    r"""Find lines in home that are inspectable.
    Walk back from the err line up to 3 lines, but don't walk back over
    changes in indent level.
    Walk forward up to 3 lines, counting \ separated lines as 1. Don't walk
    over changes in indent level (unless part of an extended line)"""
    cnt = re.compile(r'\\[\s\n]*$')
    df = re.compile(r':[\s\n]*$')
    ind = re.compile(r'^(\s*)')
    toinspect = []
    home = lines[pos]
    home_indent = ind.match(home).groups()[0]
    before = lines[max(pos-3, 0):pos]
    before.reverse()
    after = lines[pos+1:min(pos+4, len(lines))]
    for line in before:
        if ind.match(line).groups()[0] == home_indent:
            toinspect.append(line)
        else:
            break
    toinspect.reverse()
    toinspect.append(home)
    home_pos = len(toinspect)-1
    continued = cnt.search(home)
    for line in after:
        if (
            (continued or ind.match(line).groups()[0] == home_indent)
            and not df.search(line)
        ):
            toinspect.append(line)
            continued = cnt.search(line)
        else:
            break
    log.debug("Inspecting lines '''%s''' around %s", toinspect, home_pos)
    return toinspect, home_pos


class Expander:
    """Simple expression expander. Uses tokenize to find the names and
    expands any that can be looked up in the frame."""
    def __init__(self, locals, globals):
        self.locals = locals
        self.globals = globals
        self.lpos = None
        self.expanded_source = ''

    def __call__(self, ttype, tok, start, end, line):
        if self.lpos is not None:
            if start[1] >= self.lpos:
                self.expanded_source += ' ' * (start[1]-self.lpos)
            elif start[1] < self.lpos:
                self.expanded_source += ' ' * start[1]
        self.lpos = end[1]
        if ttype == tokenize.INDENT:
            pass
        elif ttype == tokenize.NAME:
            try:
                val = self.locals[tok]
                if isinstance(val, collections.abc.Callable):
                    val = tok
                else:
                    val = repr(val)
            except KeyError:
                try:
                    val = self.globals[tok]
                    if isinstance(val, collections.abc.Callable):
                        val = tok
                    else:
                        val = repr(val)

                except KeyError:
                    val = tok
            self.expanded_source += val
        else:
            self.expanded_source += tok
        if re.match(r'\s+\\\n', line[end[1]:]):
            self.expanded_source += ' \\\n'