import smp.builtins
import traceback
import inspect
import sys

from typing import Any
from enum import Enum
from io import StringIO
from contextlib import redirect_stdout


class ParserState(Enum):
    NORMAL = 1
    IN_QUOTES = 2
    IN_MACRO = 3
    IN_MACRO_ARGS = 4
    IN_SPECIAL_MACRO = 5
    IN_SPECIAL_MACRO_EXPRESSION = 6
    IN_CODE = 7
    DNL = 8


def debug_hi(input, linestart, pos, color="\u001b[7m"):
    return debug_hi_range(input, linestart, pos, pos + 1, color=color)


def debug_hi_range(input, linestart, start, end, color="\u001b[7m"):
    lineend = input.find("\n", linestart)
    line = input[linestart:] if lineend == -1 else input[linestart:lineend]

    out = ""
    for j, _c in enumerate(line):
        c_pos = linestart + j
        if (c_pos >= start) and (c_pos < end):
            out += f"{color}{_c}\u001b[0m"
        else:
            out += f"{_c}"

    return out


def macro_is_whitespace_deleting(s: str) -> bool:
    if len(s) == 0:
        return False
    return s[-1] == "_"


def macro_name_clean(macro_name: str) -> str:
    if macro_is_whitespace_deleting(macro_name):
        macro_name = macro_name[:-1]
    return macro_name


class MacroProcessor:
    source_file_path: str
    """All currently defined macros in this MacroProcessor"""
    macros: dict[str, Any]
    """ Macros which are @directives """
    special_macros: dict[str, tuple[Any, Any]]
    """ Global environment for python execution """
    py_global_env: dict
    """ Local environment for python execution """
    py_local_env_alt: dict
    py_local_env_current: dict
    """ All macro invocations that has happened """
    macro_invocations: list[tuple[str, list[str]]]
    warnings: list[Any]
    start_quote: str = '%"'
    end_quote: str = '"%'
    prefix: str = ""

    file = None

    expansion_stack: list[Any]
    file_stack: list[Any]

    def __init__(self, prefix=""):
        self.macros = dict()
        self.macro_invocations = list()
        self.warnings = list()
        self.py_global_env = dict()
        self.py_local_env_alt = dict()
        self.py_local_env_current = self.macros
        self.prefix = prefix

        self.expansion_stack = (
            list()
        )  # This probably needs some magic at some point later
        self.file_stack = list()  # This probably needs some magic at some point later

        self._define_builtins(self.macros)
        self._define_builtins(self.py_local_env_alt)

        self.special_macros = {
            "html_from_markdown": (smp.builtins.smp_builtin_html_from_markdown, None),
            "template": (smp.builtins.smp_builtin_template, None),
        }

    def _import_symbols(self, module, env, function_prefix=""):
        for name, fun in inspect.getmembers(module, inspect.isfunction):
            if name.startswith("_"):
                continue
            if not name.startswith(function_prefix):
                continue
            name = name.replace(function_prefix, "")
            env[f"{self.prefix}{name}"] = fun

    def _define_builtins(self, env):
        self._import_symbols(smp.builtins, env, function_prefix="smp_builtin_")

        env[f"{self.prefix}macro_processor"] = self
        env[f"{self.prefix}template_stack"] = []
        env[f"{self.prefix}template_stack_content"] = []
        env[f"{self.prefix}template_prefix"] = "templates/"

        # If true, include-macros will parse yaml in beginning of content
        env[f"{self.prefix}parse_file_yaml"] = True
        # If true, some macros will run in a draft-mode,
        # meaning they will skip steps that are slow.
        env[f"{self.prefix}draft"] = False
        env[f"{self.prefix}metadata_prefix"] = "METADATA_"
        env[f"{self.prefix}smp_debug"] = "4"

    def define_macro_string(self, macro_name, macro_value):
        self.define_macro(macro_name, str(macro_value))

    def define_macro(self, macro_name, macro_value):
        self.macros[macro_name] = macro_value

    def _define_macro_builtin(self, macro_name, macro_value, sub_prefix: str = ""):
        self.define_macro(f"{self.prefix}{sub_prefix}{macro_name}", macro_value)

    def _get_macro_builtin(self, macro_name, sub_prefix: str = "", default=None):
        return self.macros.get(f"{self.prefix}{sub_prefix}{macro_name}", default)

    def _get_metadata(self, macro_name, default=None):
        sub_prefix = self._get_macro_builtin("metadata_prefix")
        return self.macros.get(f"{sub_prefix}{macro_name}", default)

    def _define_metadata(self, macro_name, macro_value):
        sub_prefix = (self._get_macro_builtin("metadata_prefix"),)
        self.define_macro(f"{sub_prefix}{macro_name}", macro_value)

    def log_warning(self, message):
        """
        Here we should add some more information, line number, file etc, when that is available
        """
        self.warnings.append(message)

    def _expand_callable_macro(self, macro: Any, args: list[str] = list()) -> str:
        signature = inspect.signature(macro)
        macro_args: list[Any] = []
        if "macro_processor" in signature.parameters or "smp" in signature.parameters:
            macro_args.append(self)
        macro_args.extend(args)
        # This should maybe be processed as well(?)
        res = str(macro(*macro_args))
        return res

    def _expand_string_macro(self, macro: Any, args: list[str] = list()) -> str:
        expanded = macro
        for i, arg in enumerate(args):
            placeholder = f"${i}"
            expanded = expanded.replace(placeholder, arg)

        res = self.process_input(expanded)
        return res

    def _expand_unknown_macro(
        self, macro_name: str, args: list[str] = list(), process_args: bool = True
    ) -> str:
        if len(args) == 0:
            return macro_name
        out = f"{macro_name}("
        for i, arg in enumerate(args):
            if process_args:
                out += self.process_input(arg)
            else:
                out += arg
            if i < (len(args) - 1):
                out += ","
        out += ")"
        return out

    def expand_macro(self, macro_name: str, args: list[str] = list()) -> str:
        # Ignore trailing underscore in macro name, the parser will pop a space in front if
        # present, but we should ignore it for finding the macro.
        macro_name = macro_name_clean(macro_name)
        if macro_name not in self.macros:
            return self._expand_unknown_macro(macro_name, args, process_args=True)

        # Strip leading whitespace from arguments
        for arg in args:
            arg = arg.strip()

        # Log macro invokation
        # The fact that we are here, does not ensure that the macro is actually expanded into
        # something useful, just that it exists, and was invoked
        self.macro_invocations.append((macro_name, args))

        macro = self.macros.get(macro_name)

        if callable(macro):
            try:
                return self._expand_callable_macro(macro, args)
            except Exception as e:
                s = self._expand_unknown_macro(macro_name, args, process_args=False)

                e2 = ""
                if self._debug_on(1):
                    e2 += "\n"
                    e2 += self._stack()

                if self._debug_on(2):
                    e2 += "\n- Python stack ---\n"
                    e2 += traceback.format_exc()

                self.log_warning(f"Error expanding macro {s} ({e}){e2}")
                return s

        if isinstance(macro, str):
            return self._expand_string_macro(macro, args)

        return f"{repr(macro)}"

    def _enter_frame(self, frame_type, filename, position, content):
        self.expansion_stack.append((frame_type, filename, position, ""))

    def _pop_frame(self):
        self.expansion_stack.pop()

    def _enter_file_frame(self, filename, linenr, template):
        self.file_stack.append([filename, linenr, 0])

    def _pop_file_frame(self):
        self.file_stack.pop()

    def _stack(self):
        out = "- File stack ------\n"
        for i, frame in enumerate(self.file_stack):
            out += f" {' ' * i} {frame[0]}:{frame[1]}\n"
        out += "\n- Macro stack -----\n"
        for i, frame in enumerate(self.expansion_stack):
            out += f" {' ' * i} [{frame[0]}] in file {frame[1]}:{frame[2]}\n"
        return out

    def _debug_on(self, level):
        try:
            if level > int(self._get_macro_builtin("smp_debug")):
                return False
        except:
            return False

        return True

    def _debug(self, level, message):
        if self._debug_on(level):
            return
        print(message, file=sys.stderr)

    def process_input(self, input: str, file: str | None = None):
        """
        Extend special macro syntax to support nesting, blocks, etc
        @if <python-expression>

        @else <python-expression>

        @endif

        @for <python-expression>
        @endfor

        Note: Consider writing a new implementation that does it the same way M4 does,
        by pushing the expanded macros back to the input string, this may be more confusing,
        but may also be faster (stream or mutable string)
        """
        if file is not None:
            self.file = file
        if (file is None) and (self.file is not None):
            file = self.file
        output = ""
        state = ParserState.NORMAL
        state_start = 0
        macro_name = ""
        macro_args = []
        argument = ""
        py_expr = ""

        skip_next_line_ending = False

        # We should keep track of filename, linenumber, and character number on line here
        # So we can give sensible error messages
        # Probably add to python stack trace?
        linestart = 0
        linenr = 0

        quote_level = 0
        parens_level = 0

        i = 0
        while i < len(input):
            c = input[i]
            peek = None if i + 1 >= len(input) else input[i + 1]

            def hi_range(end=None):
                if not self._debug_on(5):
                    return
                _linestart = linestart
                _linenr = linenr
                if state_start < linestart:
                    _linestart = state_start
                    _linenr = f"~{linenr-1}"
                if end is None:
                    end = i
                range = debug_hi_range(
                    input, _linestart, state_start, end, color="\u001b[42m"
                )
                print(
                    f"[{file}:{_linenr}] {range} \t\u001b[33m({state})\u001b[0m",
                    file=sys.stderr,
                )

            if c == "\n":
                linenr += 1
                linestart = i + 1
                self.file_stack[-1][1] = linenr

            if self._debug_on(5):
                print(
                    f"[{file}:{linenr}] {debug_hi(input, linestart, i)} \t\u001b[33m({state})\u001b[0m",
                    file=sys.stderr,
                )

            if state == ParserState.DNL:
                if c == "\n":
                    state = ParserState.NORMAL
            elif state == ParserState.NORMAL:
                if skip_next_line_ending and (c == "\n"):
                    skip_next_line_ending = False
                    i += 1
                    continue

                if c == "%" and peek == "(":
                    state = ParserState.IN_CODE
                    i += 2
                    state_start = i
                    continue

                if c == "@" and i == linestart:
                    nline = input.find("\n", i)
                    nspace = input.find(" ", i)
                    if nline > -1 or nspace > 0.1:
                        args = []

                        end_special = nline
                        if -1 < nspace < nline:
                            end_special = nspace
                            args = [
                                x.strip()
                                for x in input[nspace:nline].split(" ")
                                if x is not ""
                            ]

                        sname = input[i + 1 : end_special]
                        if sname in self.special_macros:
                            hi_range(end=end_special)
                            args.append(input[nline + 1 :])
                            # @TODO: Make try/catch wrapper here as well
                            output += self._expand_callable_macro(
                                self.special_macros[sname][0], args
                            )
                            i = len(input) - 1
                            continue

                if c == "%" and peek == '"':
                    state = ParserState.IN_QUOTES
                    quote_level += 1
                    i += 1
                    state_start = i
                elif c.isalnum():
                    state = ParserState.IN_MACRO
                    state_start = i
                    macro_name += c
                else:
                    output += c

            elif state == ParserState.IN_QUOTES:
                if c == "%" and peek == '"':
                    quote_level += 1
                    i += 1
                    output += '%"'
                elif c == '"' and peek == "%":
                    quote_level -= 1
                    if quote_level == 0:
                        state = ParserState.NORMAL
                        state_start = i
                    else:
                        output += '"%'
                    i += 1
                else:
                    output += c

            elif state == ParserState.IN_MACRO:
                if c.isalnum() or c == "_":
                    macro_name += c
                elif c == "(":
                    hi_range()
                    parens_level += 1
                    state = ParserState.IN_MACRO_ARGS
                    state_start = i
                else:
                    hi_range()
                    if macro_is_whitespace_deleting(macro_name):
                        if output[-1] == " ":
                            output = output[:-1]
                        macro_name = macro_name_clean(macro_name)

                    self._enter_frame(macro_name, file, linenr, input)

                    if macro_name == "SNNL":
                        skip_next_line_ending = c != "\n"
                    elif macro_name == "DNL":
                        if c != "\n":
                            state = ParserState.DNL
                            state_start = i
                        macro_name = ""
                        i += 1
                        self._pop_frame()
                        continue
                    else:
                        expanded = self.expand_macro(macro_name)
                        output += expanded
                        output += c
                    macro_name = ""

                    self._pop_frame()

                    state = ParserState.NORMAL
                    state_start = i
            elif state == ParserState.IN_MACRO_ARGS:
                if c == "%" and peek == '"':
                    quote_level += 1
                    i += 2
                    argument += '%"'
                    continue
                elif c == '"' and peek == "%":
                    quote_level -= 1
                    i += 2
                    argument += '"%'
                    continue
                elif quote_level > 0:
                    argument += c
                    i += 1
                    continue

                if (c == ")") and (parens_level == 1):
                    hi_range()
                    if macro_is_whitespace_deleting(macro_name):
                        if output[-1] == " ":
                            output = output[:-1]
                        macro_name = macro_name_clean(macro_name)
                    parens_level = 0
                    macro_args.append(argument.strip())
                    self._enter_frame(macro_name, file, linenr, input)
                    expanded = self.expand_macro(macro_name, macro_args)
                    output += expanded
                    state = ParserState.NORMAL
                    state_start = i
                    macro_name = ""
                    macro_args = []
                    argument = ""
                    self._pop_frame()
                elif (c == ",") and (parens_level == 1):
                    macro_args.append(argument.strip())
                    state_start = i
                    hi_range()
                    argument = ""
                else:
                    if c == "(":
                        parens_level += 1
                    if c == ")":
                        parens_level -= 1
                    argument += c
            elif state == ParserState.IN_CODE:
                if c == ")" and peek == "%":
                    try:
                        self._enter_frame("inline_code", file, linenr, input)
                        f = StringIO()
                        with redirect_stdout(f):
                            exec(py_expr, self.py_global_env, self.py_local_env_current)
                        s = f.getvalue()
                        if s != "":
                            output += s
                    except Exception:
                        traceback.print_exc()
                    finally:
                        self._pop_frame()
                    py_expr = ""
                    state = ParserState.NORMAL
                    i += 1
                    state_start = i
                else:
                    py_expr += c
            i += 1

        # Handle cases where the text ends with a macro without arguments
        if macro_name != "":
            if macro_is_whitespace_deleting(macro_name):
                if len(output) > 0 and output[-1] == " ":
                    output = output[:-1]
                macro_name = macro_name_clean(macro_name)
            hi_range()
            output += self.expand_macro(macro_name)
        return output

    def store(self, **xargs):
        requested_keys = self._get_metadata("keep_states", [])
        for key in self.macros.keys():
            if (
                key.startswith(self._get_macro_builtin("metadata_prefix"))
                and key not in requested_keys
            ):
                requested_keys.append(key)

        if isinstance(requested_keys, str):
            requested_keys = [str(requested_keys)]

        needs_recompilation = (
            f"{self._get_macro_builtin('metadata_prefix')}keep_states" in self.macros
        ) or (f"{self.prefix}all_tagged_by" in [x[0] for x in self.macro_invocations])

        target_filename = self._get_metadata("target_filename")

        self.py_global_env["macro_processor_state"][self.source_file_path] = dict(
            {
                # "content": "",
                "stored_data": {
                    k: v for k, v in self.macros.items() if k in requested_keys
                },
                "extension": self._get_macro_builtin("target_file_extension"),
                "source_path": self.source_file_path,
                "needs_recompilation": needs_recompilation,
                "target_filename": target_filename,
                **xargs,
            }
        )
        return self.py_global_env["macro_processor_state"][self.source_file_path]


class MacroProcessorState:
    global_state: dict

    def __init__(self):
        self.global_state = dict()

    def macro_processor(self, macro_processor=None, prefix=""):
        if macro_processor is None:
            macro_processor = MacroProcessor(prefix)

        macro_processor.py_global_env["macro_processor_state"] = self.global_state
        return macro_processor

    def print_state(self):
        for key, val in self.global_state.items():
            print(f"{key[-20:]:20} {val}")