import os
import sys
import traceback
from typing import Any, Optional, List

from pyspark.sql.utils import AnalysisException

from spark_pipeline_framework.utilities.cannot_cast_exception_parser import parse_cannot_cast_exception


class FriendlySparkException(Exception):
    # noinspection PyUnusedLocal
    def __init__(
        self, exception: Exception, stage_name: Optional[str], *args: Any,
        **kwargs: Any
    ) -> None:
        """
        This exception wraps Spark Exceptions to extract out all the messages and show them

        :param exception:
        :param stage_name:
        :param args:
        :param kwargs:
        """
        try:
            self.exception: Exception = exception
            self.stage_name: Optional[str] = stage_name
            # Summary is a boolean argument
            # If True, it prints the exception summary
            # This way, we can avoid printing the summary all
            # the way along the exception "bubbling up"
            error_text: str = (
                stage_name or ''
            ) + ": " + FriendlySparkException.exception_summary()
            self.message: str = error_text
            if isinstance(exception, AnalysisException):
                self.message = str(exception)
            super().__init__(error_text)
        except KeyError:
            pass

    # noinspection SpellCheckingInspection
    @staticmethod
    def get_errortext(text: str) -> str:
        # Makes exception summary both BOLD and RED (FAIL)
        return text

    @staticmethod
    def exception_summary() -> str:
        # Gets the error stack
        msg = traceback.format_exc()
        try:
            # Builds the "frame" around the text
            # Gets the information about the error and makes it BOLD and RED
            info = list(
                filter(
                    lambda t: len(t) and t[0] != '\t',
                    msg.split('\n')[::-1]
                )
            )
            error = FriendlySparkException.get_errortext(
                'Error\t: {}'.format(info[0])
            )
            # Figure out where the error happened - location (file/notebook), line and function
            idx = [t.strip()[:4] for t in info].index('File')
            where = [v.strip() for v in info[idx].strip().split(',')]
            location, line, func = where[0][5:], where[1][5:], where[2][3:]
            # If it is a pyspark error, just go with it
            if 'pyspark' in error:
                new_msg = '\n{}'.format(error)
            # Otherwise, build the summary
            else:
                new_msg = '\nLocation: {}\nLine\t: {}\nFunction: {}\n{}'.format(
                    location, line, func, error
                )
            # now loop  through __cause__ to build up message including all the error messages
            type_, exc, tr = sys.exc_info()
            if exc:
                current_exception: BaseException = exc
                i: int = 0
                while hasattr(
                    current_exception, "__cause__"
                ) and current_exception.__cause__:
                    current_exception = current_exception.__cause__
                    file_name: str = ""
                    line_number: Optional[int] = None
                    if hasattr(
                        current_exception, "__traceback__"
                    ) and current_exception.__traceback__:
                        file_name = os.path.basename(
                            current_exception.__traceback__.tb_frame.f_code.
                            co_filename
                        )
                        line_number = current_exception.__traceback__.tb_frame.f_code.co_firstlineno
                    i += 1
                    exception_message: str = str(current_exception)
                    if "cannot cast " in exception_message:
                        result: List[str] = parse_cannot_cast_exception(
                            exception_message
                        )
                        for item in result:
                            new_msg += "\n" + (
                                "<" * i
                            ) + " " + item + " [" + file_name + (
                                f" ({line_number})" if line_number else ""
                            ) + "]"
                    new_msg += "\n" + (
                        "<" * i
                    ) + " " + exception_message + " [" + file_name + (
                        f" ({line_number})" if line_number else ""
                    ) + "]"

            return new_msg
        except Exception as e:
            # If we managed to raise an exception while trying to format the original exception...
            # Oh, well...
            return 'This is awkward... \n{}'.format(str(e))
