import io
import os
import traceback

def wrap_callable(callable):
    """ Wraps the given callable in order to bypass AttributeError
        handling of Jinja and provide better exception stacktraces.
    """
    def autojinja_wrapped_fcall(*args, **kwargs):
        try:
            result = callable(*args, **kwargs)
            return wrap_object_attributes(result)
        except AttributeError as e:
            message = clean_wrapped_stacktrace(e)
            raise DebugAttributeError(message) from None
        except BaseException as e:
            message = clean_wrapped_stacktrace(e)
            raise wrap_exception(e, message).with_traceback(None)
    return autojinja_wrapped_fcall

### Base exception

class CommonException(Exception):
    @staticmethod
    def from_marker(exceptionType, marker, pos, size, message):
        line = line_at_index(marker.string, pos)
        (l, c) = index_to_coordinates(marker.string, pos)
        message = f"{message}\n{line}\n{' ' * (c-1)}{'^' * size} line {l}, column {c}"
        exception = exceptionType(message)
        exception.coordinates = (l, c)
        return exception
    @staticmethod
    def from_exception(exception, marker):
        exception = prepend_jinja2_traceback(exception)
        text = format_text(marker.header)
        (l, c) = index_to_coordinates(marker.string, marker.header_open)
        stack = f"\n  During {'reinsertion' if marker.is_edit else 'generation'} of \"{marker.open} {text} {marker.close}\" at line {l}, column {c}"
        message = str(exception).lstrip('\n')
        message = f"{stack}\n{message}"
        exception = wrap_exception(exception, message)
        exception.coordinates = (l, c)
        return exception

    def __init__(self, message):
        super().__init__(message)

### Parsing exception

class ParsingException(CommonException):
    def __init__(self, message):
        super().__init__(message)

class OpenMarkerNotFoundException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(OpenMarkerNotFoundException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Couldn't find corresponding open marker \"{marker.open} {marker.close}\":")
    def __init__(self, message):
        super().__init__(message)

class CloseMarkerNotFoundException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(CloseMarkerNotFoundException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Couldn't find corresponding close marker \"{marker.close}\":")
    def __init__(self, message):
        super().__init__(message)

class EndMarkerNotFoundException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(EndMarkerNotFoundException,
                                           marker,
                                           marker.header_close - len(marker.close),
                                           len(marker.close),
                                           f"Couldn't find corresponding end marker \"{marker.open} {marker.end} {marker.close}\":")
    def __init__(self, message):
        super().__init__(message)

class RequireHeaderInlineException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(RequireHeaderInlineException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Marker can't have a multiline header:")
    def __init__(self, message):
        super().__init__(message)

class RequireHeaderMultilineException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(RequireHeaderMultilineException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Marker can't have a one line header:")
    def __init__(self, message):
        super().__init__(message)

class WrongHeaderIndentationException(ParsingException):
    @staticmethod
    def from_marker(marker, pos):
        return CommonException.from_marker(WrongHeaderIndentationException,
                                           marker,
                                           pos,
                                           len(marker.open),
                                           f"Wrong marker header indentation:")
    def __init__(self, message):
        super().__init__(message)

class RequireNewlineException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(RequireNewlineException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Marker can't start on same line as previous end marker:")
    def __init__(self, message):
        super().__init__(message)

class RequireInlineException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(RequireInlineException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Marker must start on same line as previous marker:")
    def __init__(self, message):
        super().__init__(message)

class WrongInclusionException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(WrongInclusionException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Marker has wrong inclusion regarding enclosing markers:")
    def __init__(self, message):
        super().__init__(message)

class DuplicateEditException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(DuplicateEditException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Duplicate edit marker \"{marker.open} {marker.header} {marker.close}\", consider reusing/removing duplicates:")
    def __init__(self, message):
        super().__init__(message)

class DirectlyEnclosedEditException(ParsingException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(DirectlyEnclosedEditException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Directly enclosed edit marker \"{marker.open} {marker.header} {marker.close}\", consider reusing/removing it:")
    def __init__(self, message):
        super().__init__(message)

### Generation exception

class GenerationException(CommonException):
    def __init__(self, message):
        super().__init__(message)

class RequireBodyInlineException(GenerationException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(RequireBodyInlineException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Generated body must contain only one line to be inlined:")
    def __init__(self, message):
        super().__init__(message)

class NonGeneratedEditException(GenerationException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(NonGeneratedEditException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Non-generated edit marker \"{marker.open} {marker.header} {marker.close}\", consider reusing/removing it:")
    def __init__(self, message):
        super().__init__(message)

class AlreadyGeneratedEditException(GenerationException):
    @staticmethod
    def from_marker(marker):
        return CommonException.from_marker(AlreadyGeneratedEditException,
                                           marker,
                                           marker.header_open,
                                           len(marker.open),
                                           f"Already generated edit marker \"{marker.open} {marker.header} {marker.close}\", consider fixing generation:")
    def __init__(self, message):
        super().__init__(message)

### Utils

def index_to_coordinates(string, index):
    """ Returns the corresponding tuple (line, column) of the character at the given index of the given string.
    """
    if index < 0:
        index =  index % len(string)
    sp = string[:index+1].splitlines(keepends=True)
    return len(sp), len(sp[-1])

def line_at_index(string, index):
    """ Returns the line at the given index of the given string.
    """
    if index < 0:
        index =  index % len(string)
    start = string.rfind('\n', 0, index)+1
    end = string.find('\n', index)
    suffix = '\\n' if end >= 0 else '\\0'
    if end < 0:
        end = len(string)
    return f"{string[start:end]}{suffix}"

def format_text(text, max_length = 50, suffix = "..."):
    """ Replaces tabulations and ends of line.
    """
    text = text if len(text) <= max_length else f"{text[:max_length]}{suffix}"
    return text.replace('\t', "\\t").replace('\n', "\\n")

def split_traceback(tb, remove_class_name = False):
    """ Splits the given traceback into (stacktrace, message).
    """
    startidx = -1
    while True:
        startidx = tb.find('\n', startidx+1)
        if startidx < 0:
            break # No stacktrace
        if not tb.startswith('  ', startidx+1):
            break # End of stacktrace
    if startidx < 0:
        stacktrace = None # No stacktrace
    else:
        stacktrace = tb[:startidx]
    if not remove_class_name:
        message = tb[startidx+1:]
    else:
        idx = tb.find(' ', startidx+1)
        if idx < 0:
            message = None # No message
        else:
            idx += 2 if tb[idx+1] == '\n' else 1 # lstrip one
            message = tb[idx:]
    return stacktrace, message

_disallowed_attrs = [
    "__delattr__",
    "__getattr__",
    "__getattribute__",
    "__init__",
    "__new__",
    "__setattr__",
    "__str__",
]

def wrap_exception(exception, message):
    """ Wraps the given exception with a custom message.
        Creates an object that mocks the exception class.
    """
    if hasattr(exception, "_wrapped_exception") and hasattr(exception, "_wrapped_message"):
        exception._wrapped_message = message
        return exception
    try:
        def __init__(self):
            pass
        def __str__(self):
            return self._wrapped_message
        dicts = { name: getattr(exception, name) for name in dir(exception)
                  if name.startswith("__") and name not in _disallowed_attrs }
        dicts["__init__"] = __init__
        dicts["__module__"] = exception.__class__.__module__
        dicts["__str__"] = __str__
        mock = type(exception.__class__.__qualname__, (exception.__class__,), dicts)()
        for key, value in mock.__dict__.items():
            setattr(mock, key, value)
        mock._wrapped_exception = exception
        mock._wrapped_message = message
        return mock
    except:
        return exception

def prepend_jinja2_traceback(exception, tb = None):
    """ Prepends the Jinja2 traceback to the given exception.
        Creates an object that derives the exception class.
    """
    if tb == None:
        tb = traceback.format_exc().rstrip('\n')
    (stacktrace, message) = split_traceback(tb, True)
    if stacktrace != None:
        token1 = f"jinja2{os.sep}environment.py\", " # Jinja2 hack
        token2 = f"jinja2{os.sep}runtime.py\", " # Jinja2 hack
        startidx = -1
        while True:
            idx = stacktrace.find(token1, startidx+1)
            if idx < 0:
                idx = stacktrace.find(token2, startidx+1)
                if idx < 0:
                    break # Token not found
                startidx = stacktrace.find('\n', idx+len(token2))
            else:
                startidx = stacktrace.find('\n', idx+len(token1))
            if startidx < 0:
                startidx = len(stacktrace)
                break # No eol after token
            if not stacktrace.startswith("    ", startidx+1):
                continue # No inner traceback
            startidx = stacktrace.find('\n', startidx+1)
            if startidx < 0:
                startidx = len(stacktrace)
                break # No eol after inner traceback
        if startidx < 0:
            return exception
        stacktrace = stacktrace[startidx+1:]
        if len(stacktrace) == 0:
            stacktrace = None
        else:
            # Clean autojinja stacktrace
            token = f"autojinja{os.sep}templates.py\", "
            idx = stacktrace.find(token)
            if idx >= 0:
                startidx = stacktrace.rfind('\n', 0, idx)
                if startidx < 0: # No eol before token
                    stacktrace = None
                else:
                    stacktrace = stacktrace[:startidx]
    if message == None:
        message = ""
    message = f"\n{stacktrace}\n{message}" if stacktrace else f"\n{message}"
    return wrap_exception(exception, message)

def clean_traceback(exception):
    """ Removes the traceback of the given exception.
    """
    return prepend_jinja2_traceback(exception).with_traceback(None)

### --debug option enabled

class DebugAttributeError(BaseException):
    def __init__(self, message):
        super().__init__(message)

_disallowed_stacktraces = [
    "autojinja_wrapped_fcall",
    "autojinja_fcall",
    "autojinja_fget",
    "autojinja_fset",
    "autojinja_fdel"
]

def clean_wrapped_stacktrace(exception):
    """ Cleans the the given exception's stacktrace from useless wrapper calls.
    """
    if exception.__class__.__module__ != "builtins":
        fullname = f"{exception.__class__.__module__}.{exception.__class__.__qualname__}: "
    else:
        fullname = f"{exception.__class__.__qualname__}: "
    stacktrace = traceback.format_exc()
    lines = stacktrace.splitlines()
    start = -1
    for i in range(0, len(lines), 1):
        if lines[i].endswith("line 15, in autojinja_wrapped_fcall") \
        or lines[i].endswith("line 18, in autojinja_wrapped_fcall"): # /!\ easily breakable
            start = i
            break
    end = -1
    if start > 0:
        for i in range(len(lines)-1, start, -1):
            if lines[i] == fullname:
                end = i+1
                break
    message = io.StringIO()
    state = 0 # 0, 1
    for line in lines[:start]:
        if state == 1 and line.startswith("    "):
            state = 0 # Skip
        elif line.startswith("Traceback (most recent call last)"):
            pass # Skip
        elif not any(x in line for x in _disallowed_stacktraces):
            message.write(f"\n{line}")
        else:
            state = 1 # Skip
    for line in lines[end:]:
        message.write(f"\n{line}")
    return message.getvalue()

_wrapped_properties = set()

def _autojinja_fcall(wrapped_callable):
    def autojinja_fcall(*args, **kwargs):
        return wrapped_callable(*args, **kwargs)
    return wrap_callable(autojinja_fcall)

def _autojinja_fget(wrapped_property):
    def autojinja_fget(self):
        return wrapped_property.fget(self)
    return wrap_callable(autojinja_fget)

def _autojinja_fset(wrapped_property):
    def autojinja_fset(self, value):
        return wrapped_property.fset(self, value)
    return wrap_callable(autojinja_fset)

def _autojinja_fdel(wrapped_property):
    def autojinja_fdel(self):
        return wrapped_property.fdel(self)
    return wrap_callable(autojinja_fdel)

_disallowed_modules = set([
    "autojinja.defaults",
    "autojinja.exceptions",
    "autojinja.main",
    "autojinja.parser",
    "autojinja.path",
    "autojinja.templates",
    "autojinja.utils",
    "builtins",
])

def wrap_object_attributes(obj):
    """ Wraps all properties and methods of the given objects in order to bypass
        AttributeError handling of Jinja and provide better exception stacktraces.
    """
    if str(obj.__class__.__module__) in _disallowed_modules:
        return obj
    for name in dir(obj):
        try:
            if not hasattr(obj.__class__, name):
                continue
            attribute = getattr(obj.__class__, name)
            if callable(attribute):
                wrapped_callable = getattr(obj, name)
                if str(wrapped_callable.__name__) == "autojinja_wrapped_fcall":
                    continue # Already replaced
                setattr(obj, name, _autojinja_fcall(wrapped_callable))
            elif isinstance(attribute, property):
                wrapped_property = attribute
                key = f"{obj.__class__}.{name}"
                if key in _wrapped_properties:
                    continue # Already replaced
                _wrapped_properties.add(key)
                setattr(obj.__class__, name, property(fget=_autojinja_fget(wrapped_property),
                                                      fset=_autojinja_fset(wrapped_property),
                                                      fdel=_autojinja_fdel(wrapped_property)))
        except:
            pass # Skip attribute
    return obj
