__author__ = "Duc Tin"

import re
import logging
import tokenize
import linecache
from io import StringIO
from contextlib import contextmanager
from traceback import StackSummary, walk_tb

ID_PATTERN = re.compile(r'([a-zA-Z_][a-zA-Z_0-9.]*)')
MEM_PATTERN = re.compile(r'<.*object at 0x[0-9A-Fa-f]+>')


def get_recur_attr(obj, attr: str):
    """
    Follow the dot `.` in attribute string `attr` to the final object
    :param obj: any object
    :param attr: string of attribute access
    :return: the final desire object or '!!! Not Exists'

    example: if we need to access `c` object as in `a.b.c`,
    then `obj=a, attr='b.c'`
    If `b` or `c` doesn't exists, then `'!!! Not Exists'` is returned
    """
    this_level, *next_levels = attr.split('.', maxsplit=1)
    if not next_levels:
        try:
            return getattr(obj, this_level)
        except AttributeError:
            return "!!! Not Exists"
        except Exception as e:
            return f"!!! Attribute Error: {e}"
    else:
        try:
            obj = getattr(obj, this_level)
        except AttributeError:
            return "!!! Not Exists"
        return get_recur_attr(obj, next_levels[0])


def get_repr(obj, multiline_indent: int = 0) -> str:
    """
    Pick a more useful representation of object `obj` between __str__ and __repr__
     if it available.
    :param obj: any of object
    :param multiline_indent: indent level for the second line onward
    :return: string of object representation
    """
    try:
        _repr_ = repr(obj)
        if not MEM_PATTERN.match(_repr_):
            return _repr_
    except Exception as e:
        _repr_ = f"!!! repr error: {e}"

    try:
        _str_ = str(obj).replace('\r\n', '\n')
        if '\n' in _str_:
            indent = ' ' * multiline_indent
            _str_ = _str_.replace('\n', '\n' + indent)

        return _str_

    except Exception as e:
        return _repr_


def is_full_statement(*lines: str) -> bool:
    """
    Check if a set of lines makes up a full python statement
    :param lines: list of line of python code
    :return: True if line in `lines` made up a full python statement
    """
    try:
        stream = StringIO()
        stream.write('\n'.join(lines))
        stream.seek(0)
        for t in tokenize.generate_tokens(stream.readline):
            pass
    except tokenize.TokenError as e:
        if 'EOF in multi-line statement' in str(e):
            return False
        raise
    else:
        return True


def get_full_statement(filename, lineno: int) -> list:
    """
    Get all lines of python file `filename` that makes up a full statement starting from `lineno`
    :param filename: path to python source code file
    :param lineno: line number at with the statement started
    :return: list of all lines that made up a python statement with indent striped
    """
    lines = []
    while True:
        line = linecache.getline(filename, lineno)
        if not line:
            break

        lines.append(line)
        if is_full_statement(*lines):
            break

        lineno += 1

    lines[-1] = lines[-1].strip('\n')  # remove newline in last line
    indent = re.search(r'^(\s+)', lines[0])
    if indent:
        indent = indent.group(1)
        return [x.replace(indent, '', 1) for x in lines]
    else:
        return lines


@contextmanager
def logging_disabled():
    logging.disable(logging.CRITICAL)
    yield
    logging.disable(logging.NOTSET)


def get_traceback_depth(trace_back) -> int:
    count = 0
    tb = trace_back
    while tb is not None:
        tb = tb.tb_next
        count += 1

    return count


def analyze_frame(trace_back, full_context: int = False) -> str:
    """
    Read out variables' content surrounding the error line of code
    :param trace_back: A traceback object when exception occur
    :param full_context: Also export local variables that is not in the error line
            full_context == 0: export only variables that appear on the error line
            full_context == 1: export variables within the error function's scope
            full_context >= 2: export variables along the function's call stack up to `full_context` level
    :return: string of analyzed frame
    """
    result = []
    full_context = max(0, int(full_context))
    # todo: add color
    bullet_1 = '|->'
    bullet_2 = '=>'
    multi_line_indent1 = 8 + len(bullet_1)
    multi_line_indent2 = 8 + len(bullet_2)
    stack_depth = get_traceback_depth(trace_back)

    with logging_disabled():
        for idx, obj in enumerate(walk_tb(trace_back)):
            frame, _ = obj

            global_var = frame.f_globals
            local_var = frame.f_locals

            summary = StackSummary.extract([obj], capture_locals=True)[0]
            line = summary.line
            if not is_full_statement(summary.line):
                line = get_full_statement(summary.filename, summary.lineno)
                line = '    '.join(line)

            txt = [f'  File "{summary.filename}", line {summary.lineno}, in {summary.name}',
                   f'    {line}']

            # todo: dump all level to different file?
            parse_level = max(full_context - 1, 0)
            if idx + 1 < (stack_depth - parse_level):
                # don't parse variables for top levels
                txt.append('')
                result.append('\n'.join(txt))
                continue

            # get value of variables on the error line
            identifiers = ID_PATTERN.findall(line)
            seen = set()
            outer = "(outer) " if idx else ""  # ground level variables are not outer for sure
            for i in identifiers:
                if i in seen or i.endswith('.'):
                    continue

                seen.add(i)
                spaces = multi_line_indent1 + len(i)
                if i in local_var:
                    value = get_repr(local_var[i], spaces)
                    txt.append(f'     {bullet_1} {i} = {value}')
                elif i in global_var:
                    spaces += len(outer)
                    value = get_repr(global_var[i], spaces)
                    txt.append(f'     {bullet_1} {outer}{i} = {value}')
                elif '.' in i:
                    # class attribute access
                    spaces += len(outer)
                    instance = i.split('.')[0]
                    obj = local_var.get(instance, global_var.get(instance))
                    attribute = get_recur_attr(obj, i[len(instance) + 1:])
                    value = get_repr(attribute, spaces)
                    scope = outer if instance in global_var else ''
                    txt.append(f'     {bullet_1} {scope}{i} = {value}')
                else:
                    # reserved Keyword or non-identifier, eg. word inside the string
                    pass

            # get value of other variables within local scope
            if full_context:
                other_local_var = set(local_var) - set(identifiers)
                if other_local_var:
                    spaces = multi_line_indent2
                    txt.extend([f'     {bullet_2} {k} = {get_repr(v, spaces + len(k))}'
                                for k, v in local_var.items() if k in other_local_var])

            txt.append('')
            result.append('\n'.join(txt))

    return '\n'.join(result)
