__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_half_ended(line: str) -> bool:
    """
    Check if we a given a last line of a multiline statement with brackets
    """
    if line.count(')') > line.count('('):
        return True
    if line.count('}') > line.count('{'):
        return True
    if line.count(']') > line.count('['):
        return True


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' in str(e):
            return False
        else:
            return True
    else:
        return True


def get_statement_up(filename: str, lineno: int) -> list:
    """The python interpreter return the last line of exception lines
        This behavior is of python 3.6, 3.7
        We have to go up to grab the full statement

        Grab maximum 10 lines
    """
    startno = lineno
    lines = []

    while lineno >= max(0, startno - 9):
        line = linecache.getline(filename, lineno)
        lines.insert(0, line)

        temp = ''.join(lines)
        if not is_half_ended(temp):
            break

        lineno -= 1

    return lines


def get_statement_down(filename: str, lineno: int) -> list:
    """The python interpreter return the first line of exception lines
        This behavior is of python 3.9
        We go down normally to grab the full statement

        Grab maximum 10 lines
    """
    lines = []
    startno = lineno
    while lineno <= startno + 9:
        line = linecache.getline(filename, lineno)
        if not line:
            break

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

        lineno += 1

    return lines


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 maximum 5 lines that made up a python statement with indent striped,
    """

    line = linecache.getline(filename, lineno)
    if is_half_ended(line):
        lines = get_statement_up(filename, lineno)
    else:
        lines = get_statement_down(filename, lineno)

    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)
