import os.path
import random
import re
import traceback
import zlib
from importlib import import_module, metadata
from os.path import dirname, basename, join
from pathlib import Path
from types import ModuleType
from typing import Optional, Union, Callable, Iterable, Dict, Tuple, List, TypedDict, Any

from ptyx.pretty_print import pretty_box, yellow

from ptyx.errors import PtyxRuntimeError, PythonBlockError, PythonExpressionError, PtyxExtensionNotFound
from ptyx.extensions import CompilerExtension

import ptyx.randfunc as randfunc
from ptyx import __version__
from ptyx.config import param
from ptyx.context import GLOBAL_CONTEXT
from ptyx.internal_types import NiceOp, PickItemAction, EvalFlags, PtyxTraceback
from ptyx.sys_info import SYMPY_AVAILABLE

from ptyx.syntax_tree import Node, SyntaxTreeGenerator, Tag, TagSyntax
from ptyx.utilities import advanced_split, numbers_to_floats, _float_me_if_you_can, latex_verbatim


class State(TypedDict, total=False):
    syntax_tree: Node
    seed: Optional[int]
    input: str
    path: Optional[Path]
    loaded_extensions: dict[str, ModuleType]
    plain_ptyx_code: str
    after_include: Optional[str]


class LastEvaluatedExpressionInfo:
    position: int | None = None
    result: str | None = None


# =============================================================================
# #ASSERT{}
# #CASE{int}...#CASE{int}...#END
# #COMMENT...#END
# #DEBUG
# #IF{bool}...#ELIF{bool}...#ELSE...#END
# #IFNUM{int}{...}
# #MACRO{name}
# #NEW_MACRO{name}...#END
# #PICK[option,int]{list}
# #PYTHON...#END
# #RAND[option,int]{list}
# #SEED{int}
# #SHUFFLE #ITEM...#ITEM...#END
# #SIGN
# #SYMPY...#END
# #TEST{bool}{...}{...}
# #+, #-,  #*, #=, #? (special operators)
# #varname or #[option,int]varname
# ## options : sympy, python, float
# ## int : round
#
# =============================================================================


# noinspection PyPep8Naming,PyMethodMayBeStatic
class LatexGenerator:
    """Convert text containing ptyx tags to plain LaTeX."""

    convert_tags = {
        "+": "ADD",
        "-": "SUB",
        "*": "MUL",
        "=": "EQUAL",
        "?": "SIGN",
        # "#": "SHARP",
    }

    # noinspection RegExpRedundantEscape
    re_varname = re.compile(r"[A-Za-z_]\w*(\[.+\])?$")
    flags: EvalFlags
    macros: dict[str, list[str | Node]]
    # TODO: use a TypedDict for context. (Extensions should use TypedDict inheriting from this one.)
    #  Since typing all of global context isn't realistic (in includes sympy!), one should
    #  split context dictionary, and **create a specific dictionary for environment variables**,
    #  like PTYX_NUM, or PTYX_LATEX.
    #  This environment variables dict should use a TypedDict for local context.
    #  Using a dataclass is also possible, but it may prove complicated for extensions,
    #  with dynamic heritage needed, so it's probably not worth it.
    context: dict[str, Any]
    backups: list[list[str]]
    # Remember the position of the last text generated by an evaluated expression, like `#{a+7}`.
    last_eval_position: LastEvaluatedExpressionInfo

    def __init__(self, compiler=None):
        self.clear()
        # Some pTyX code may be generated dynamically inside a #PYTHON code
        # block for example. So, LatexGenerator need his own STG (Syntax Tree Generator)
        # to be able to scan code at runtime.
        # (Every run implies generating a syntax tree again for each
        # of these pieces of embedded pTyX code.)
        self.parser = SyntaxTreeGenerator()
        # Access to compiler is needed by some extensions.
        self.compiler = compiler
        # Add ability to share values between compilations (needed by some
        # extensions).
        self.cache = {}

    def clear(self, main_file: Path | None = None) -> None:
        self.macros = {}
        self.context = GLOBAL_CONTEXT.copy()
        # PTYX_LATEX will contain a list of generated LaTeX code parts, to be eventually concatenated.
        self.context["PTYX_LATEX"]: list[str] = []  # type: ignore
        # This is used to keep a track of included files, to be able to generate a traceback when debugging.
        self.context["PTYX_TRACEBACK"]: PtyxTraceback = [(str(main_file), None)]  # type: ignore
        self.backups = []
        self.last_eval_position = LastEvaluatedExpressionInfo()
        # When write() is called from inside a #PYTHON ... #END code block,
        # its argument may contain pTyX code needing parsing.
        # Write argument used to be parsed by default, but this was rarely needed
        # and led to very surprising behaviour sometimes from user perspective.
        # Now, one should use write(..., parse=True) if needed, which is much saner.
        self.context["write"] = self.write
        # Internal flags
        self.flags = EvalFlags(eval_as_float=self.context.get("ALL_FLOATS", False))

    def reset(self):
        """To overwrite."""
        self.clear()

    def set_new_context(self, base_context: Optional[Dict] = None) -> dict[str, Any]:
        """Set a new context of evaluation for code, except for PTYX_* variables."""
        old_context = self.context
        if base_context is None:
            base_context = GLOBAL_CONTEXT
        # Use deepcopy instead ?
        new_context = base_context.copy()
        # Copy internal parameters to new context.
        for key in old_context:
            if key.startswith("PTYX_"):
                new_context[key] = old_context[key]
        self.context = new_context
        return old_context

    @property
    def NUM(self):
        return self.context["PTYX_NUM"]

    @property
    def WITH_ANSWERS(self):
        return self.context.get("PTYX_WITH_ANSWERS")

    def _open_temp_context(self):
        # Store generated text in a temporary location, instead of self.context['LATEX'].
        # Function `function` will then be applied to this text,
        # before text is appended to self.context['LATEX'].
        # Backups may need to be read when parsing #= special tag,
        # so store them in `self.backups`.
        self.backups.append(self.context["PTYX_LATEX"])
        self.context["PTYX_LATEX"] = []

    def _apply_func_and_close_temp_context(self, function, **options):
        code = "".join(self.context["PTYX_LATEX"])
        if callable(function):
            function = [function]
        for f in function:
            code = f(code, **options)
        self.context["PTYX_LATEX"] = self.backups.pop()
        self.write(code)

    def parse_node(
        self,
        node: Node,
        function: Optional[Union[Callable, Iterable[Callable]]] = None,
        **options,
    ):
        """Parse a node in a pTyX syntax tree.

        Return True if block content was recursively parsed, and False else.

        In particular, when parsing an IF block, True will be returned if and
        only if the block condition was satisfied.
        """
        if function is not None:
            self._open_temp_context()
        name = node.name
        assert isinstance(name, str), repr(node)
        tag = self.convert_tags.get(name, name)
        try:
            method = getattr(self, f"_parse_{tag}_tag")
            return method(node)
        except AttributeError:
            print(f"Error: method '_parse_{tag}_tag' not found.")
            raise
        except Exception:
            print(f"Error when calling method '_parse_{tag}_tag'.")
            raise
        finally:
            if function is not None:
                self._apply_func_and_close_temp_context(function, **options)

    def _parse_children(
        self,
        children: Iterable[Union[str, Node]],
        function: Optional[Union[Callable, Iterable[Callable]]] = None,
        **options,
    ) -> None:
        """Parse all children nodes.

        Resulting LaTeX code will be appended to `self.context['PTYX_LATEX']`.

        If `function` is not None, apply `function` to resulting LaTeX
        code before appending it to `self.context['PTYX_LATEX']`.
        (So, `function` signature must be: function(str, **options) -> str).

        Note that `function` can also be a list of functions.
        In that case, the first function of the list is also the first
        to be applied.
        """

        if function is not None:
            self._open_temp_context()

        for child in children:
            if isinstance(child, str):
                self.write(child)
            else:
                # All remaining children should be nodes now.
                assert isinstance(child, Node)
                # Nodes are either numbered, or have a name.
                # Numbered nodes correspond to command arguments. Those should
                # have been processed before, and not be passed to _parse_children().
                assert isinstance(child.name, str), (
                    f"Argument {child.name!r} should have been processed "
                    "and removed before calling _parse_children() !"
                )

                self.parse_node(child)

        if function is not None:
            self._apply_func_and_close_temp_context(function, **options)

    @staticmethod
    def _parse_options(node: Node):
        """Parse a tag options, following the syntax {key1=val1,...}."""
        options = node.options
        args = []
        kw = {}
        if options is not None:
            options_list = options.split(",")
            for option in options_list:
                if "=" in option:
                    key, val = option.split("=")
                    kw[key.strip()] = val.strip()
                else:
                    args.append(option.strip())
        return args, kw

    def write(self, text: str, parse: bool = False, verbatim: bool = False) -> None:
        """Append a piece of LaTeX text to context['PTYX_LATEX'].

        :param text: a block of text, which may contain pTyX code.
        :type text: string
        :param parse: indicate text have to be parsed.
        :type parse: bool
        :param verbatim: emulate LaTeX verbatim mode (escape LaTeX characters, keep indentation...)
        :type verbatim: bool

        .. note:: Param `parse` defaults to False, since in most cases text is already
                  parsed at this state.
                  As an exception, when called from inside a #PYTHON [...] #END
                  code block, write() may be applied to unparsed text.
        """
        # Parsing at this stage is quite expansive yet unnecessary most of the time,
        # so some basic testing is done before.
        text = str(text)
        if text.strip():
            match (flags := self.flags).previous_nice_op:
                case NiceOp.ADD:
                    text = "+" + text
                    flags.previous_nice_op = NiceOp.NONE
                case NiceOp.SUB:
                    text = "-" + text
                    flags.previous_nice_op = NiceOp.NONE
                case NiceOp.MUL:
                    text = r"\times " + text
                    flags.previous_nice_op = NiceOp.NONE

        if parse and "#" in text:
            if param["debug"]:
                print("Parsing %s..." % repr(text))
            self.parse_node(self.parser.generate_tree(text), function=(latex_verbatim if verbatim else None))
        else:
            if verbatim:
                text = latex_verbatim(text)
            self.context["PTYX_LATEX"].append(text)

    def read(self) -> str:
        return "".join(self.context["PTYX_LATEX"])

    def _parse_APART_tag(self, node: Node):
        """Interpret a piece of code in a sandbox, eliminating side effects."""
        # Backup local variables and reset context.
        context_backup = self.context
        self.set_new_context()
        # Interpret code
        self._parse_children(node.children)
        # Restore local variables and update generated LaTeX code.
        self.set_new_context(context_backup)

    @staticmethod
    def _prettify_path(path: Path | str | None = None) -> str:
        if path is None:
            return "?"
        path = Path(path)
        return str(path.parent / f"\u001b[36m{path.name}\u001b[0m").replace("#", "##")

    def _parse_INCLUDE_START_tag(self, node: Node):
        # The file being included:
        file_path = node.arg(0)
        # The position (in the raw parent file) where the included file was inserted:
        position = node.arg(1)
        self.context["PTYX_TRACEBACK"].append((file_path, int(position) if position else None))
        print(f"\u001b[36mIMPORTING\u001b[0m '{self._prettify_path(file_path)}'")

    def _parse_INCLUDE_END_tag(self, node: Node):
        self.context["PTYX_TRACEBACK"].pop()

    def _parse_PTYX_VERSION_tag(self, node: Node):
        def version_tuple(version_: str):
            return version_.split(".")

        assert isinstance(node.children[0], Node), repr(node)
        assert isinstance(node.children[0].children[0], str), repr(node)
        version = version_tuple(node.children[0].children[0])
        if version_tuple(__version__) < version:
            print("Warning: pTyX engine is too old (v%s required)." % version)
        self.context["PTYX_VERSION"] = version

    def _parse_ASK_tag(self, node: Node):
        self._parse_children(node.children, function=self.context.get("format_ask"))

    def _parse_ASK_ONLY_tag(self, node: Node):
        if not self.WITH_ANSWERS:
            self._parse_children(node.children, function=self.context.get("format_ask_only"))
        else:
            print("Skipping ASK_ONLY section...")

    def _parse_ANS_tag(self, node: Node):
        if self.WITH_ANSWERS:
            self._parse_children(node.children, function=self.context.get("format_ans"))

    def _parse_ANSWER_tag(self, node: Node):
        if self.WITH_ANSWERS:
            assert isinstance(node.children[0], Node), repr(node)
            self._parse_children(node.children[0].children, function=self.context.get("format_answer"))

    def _parse_QUESTION_tag(self, node: Node):
        if not self.WITH_ANSWERS:
            assert isinstance(node.children[0], Node), repr(node)
            self._parse_children(node.children[0].children)

    def _parse_IF_tag(self, node: Node):
        test = node.eval_arg(0, self.context)
        if test:
            self._parse_children(node.children[1:])
        return test

    _parse_ELIF_tag = _parse_IF_tag

    def _parse_CONDITIONAL_BLOCK_tag(self, node: Node):
        for child in node.children:
            assert isinstance(child, Node)
            if self.parse_node(child):
                # If an `IF` or `ELIF` node was processed, all successive `ELIF`
                # or `ELSE` nodes must be skipped.
                # (The same for `CASE`).
                break

    def _parse_ELSE_tag(self, node: Node):
        self._parse_children(node.children)

    def _parse_IFNUM_tag(self, node: Node):
        if node.eval_arg(0, self.context) == self.NUM:
            assert isinstance(node.children[1], Node), repr(node)
            self._parse_children(node.children[1].children)

    def _parse_CASE_tag(self, node: Node):
        test = node.eval_arg(0, self.context) == self.NUM
        if test:
            self._parse_children(node.children[1:])
        return test

    def _parse_IMPORT_tag(self, node: Node):
        exec("from %s import *" % node.arg(0), self.context)

    def _parse_LOAD_tag(self, node: Node):
        # LOAD tag is used to load extensions **before** syntax tree is built.
        pass

    def _parse_INCLUDE_tag(self, node: Node):
        # INCLUDE tag is used to insert code in isolated mode **before** syntax tree is built.
        pass

    def _parse_PYTHON_tag(self, node: Node):
        assert len(node.children) == 1
        python_code = node.children[0]
        assert isinstance(python_code, str), repr(python_code)
        self._exec_python_code(python_code, self.context)

    # Remove comments before generating tree ?
    def _parse_COMMENT_tag(self, node: Node):
        pass

    def _parse_ASSERT_tag(self, node: Node):
        test = node.eval_arg(0, self.context)
        if not test:
            print("Error in assertion (NUM=%s):" % self.NUM)
            print("***")
            print(node.arg(0))
            print("***")
            assert test

    def _parse_EVAL_tag(self, node: Node):
        if self.flags.suppress_next_eval:
            self.flags = EvalFlags(eval_as_float=self.context.get("ALL_FLOATS", False))
            return
        args, kw = self._parse_options(node)
        for arg in args:
            if arg.isdigit():
                self.flags.round = int(arg)
            elif arg == ".":
                self.flags.keep_dot_as_decimal_mark = True
            elif arg == "*":
                self.flags.is_mul_coeff = True
            elif arg in ("floats", "float"):
                self.flags.eval_as_float = True
            elif arg == "str":
                self.flags.format_as_str = True
            elif arg == "select":
                self.flags.pick_action = PickItemAction.SELECT_FROM_NUM
            elif arg == "rand":
                self.flags.pick_action = PickItemAction.RAND_CHOICE
            else:
                raise PtyxRuntimeError(f"Unknown flag {arg!r} in `#[{node.options}]{node.arg(0)}`.")
        code = node.arg(0)
        assert isinstance(code, str), type(code)
        tag, code = self._get_python_code_tag(code)
        try:
            txt = self._eval_and_format_python_expr(code)
        except Exception as e:
            # print("ERROR: Can't evaluate this: " + repr(code))
            raise PythonExpressionError(
                python_code=code, flags=node.options, label=tag, context=self.context
            ) from e
        # Tags must be cleared *before* calling .write(txt), since .write(txt)
        # add '+', '-' and '\times ' before txt if corresponding flags are set,
        # and ._eval_and_format_python_expr() has already do this.
        self.flags = EvalFlags(
            eval_as_float=self.context.get("ALL_FLOATS", False),
            suppress_next_eval=self.flags.suppress_next_eval,
        )
        self.last_eval_position.position = len(self.context["PTYX_LATEX"])
        self.last_eval_position.result = txt
        self.write(txt)

    def _parse_MACRO_tag(self, node: Node):
        name = node.arg(0)
        self.macros[name] = node.children[1:]

    def _parse_CALL_tag(self, node: Node):
        """Calling a macro."""
        name = node.arg(0)
        if name not in self.macros:
            raise NameError(f"Error: MACRO {name!r} undefined.")
        self._parse_children(self.macros[name])

    def _parse_ENUM_tag(self, node: Node):
        # This tag does nothing by itself, but is used by some extensions.
        self._parse_children(node.children)

    def _parse_VERBATIM_tag(self, node: Node):
        self._parse_children(node.children, function=latex_verbatim)

    @staticmethod
    def _child_index(children, name):
        """Return the index of the first child of name `name` in `children` list.

        Argument `name` must be a `Node` name.

        A `ValueError` is raised if no such `Node` is found.
        """
        for i, child in enumerate(children):
            if isinstance(child, Node) and child.name == name:
                return i
        raise ValueError(f"No {name} Node found.")

    def _shuffle_and_parse_children(self, node: Node, children=None, target: str = "ITEM", **kw):
        # Shuffles all the #ITEM sections inside a #SHUFFLE block.
        # Note that they may be some text or nodes before first #ITEM,
        # if so they should be left unmodified at first position.
        # Nota: children may be only some children of given node.
        # This is used when some children were already processed before
        # (this is the case for arguments usually).
        if children is None:
            children = node.children
        if not children:
            return
        try:
            i = self._child_index(children, target)
        except ValueError:
            # `target` not found: nothing to shuffle.
            print(f"WARNING: Nothing to shuffle ({target!r} not found).")
            self._parse_children(children, **kw)
            return
        items = children[i:]
        # Now, we will shuffle only items which are nodes of name `target`.
        # So, we split items in groups, so that every group starts with
        # a Node named `target`.
        # For example, if target is `QUESTION`, and items are
        # ['QUESTION', 'OTHER', 'QUESTION', 'OTHER', 'ANOTHER'],
        # groups will be [['QUESTION', 'OTHER'],
        #                 ['QUESTION', 'OTHER', 'ANOTHER']]
        # Then, we will shuffle those groups. We may obtain for example :
        #                [['QUESTION', 'OTHER', 'ANOTHER'],
        #                 ['QUESTION', 'OTHER']]
        # Finally, we flatten `groups` list to obtain the following:
        # ['QUESTION', 'OTHER', 'ANOTHER', 'QUESTION', 'OTHER']
        # 1. Generate groups
        groups: list[list[Node | str]] = []
        for item in items:
            if isinstance(item, Node) and item.name == target:
                groups.append([item])
            else:
                groups[-1].append(item)
        # 2. Shuffle groups
        randfunc.shuffle(groups)
        #        print(groups)
        #        input('-- pause --')
        # 3. Flatten `groups` list
        items = sum(groups, [])

        #        print(f'items: {items}')
        #        input('-- pause --')
        self._parse_children(children[:i] + items, **kw)
        # print('\n------------')
        # print('SHUFFLE: %s elements, excluding first %s.' % len(children), i)
        # print('state hash is %s' % hash(random.getstate()))
        # print('------------\n')

    def _parse_SHUFFLE_tag(self, node: Node) -> None:
        self._shuffle_and_parse_children(node)

    def _parse_ITEM_tag(self, node: Node) -> None:
        self._parse_children(node.children)

    def _parse_SEED_tag(self, node: Node):
        # SEED tag is a special tag which is managed independently, at first pass
        # (it avoids user including it inadvertently inside
        # a conditional block or a #ASK_ONLY/#END block, which may result in
        # what seems to be a very strange behaviour of the compiler !)
        pass

    def _pick_and_parse_children(self, node: Node, children: List = None, target: str = "ITEM", **kw):
        # Choose only one between all the #ITEM sections inside a #PICK block.
        # Note that they may be some text or nodes before first #ITEM,
        # if so they should be left unmodified at their original position.
        if children is None:
            children = node.children
        if not children:
            return
        i = self._child_index(children, target)
        items = children[i:]
        for item in items:
            if not (isinstance(item, Node) and item.name == target):
                log = [
                    "This is current structure:",
                    node.display(),
                    rf"\n{item!r} is not an {target!r} node !",
                ]
                raise PtyxRuntimeError("\n".join(log))
        item = randfunc.randchoice(items)
        self._parse_children(children[:i] + [item], **kw)
        # print('\n------------')
        # print('SHUFFLE: %s elements, excluding first %s.' % len(children), i)
        # print('state hash is %s' % hash(random.getstate()))
        # print('------------\n')

    def _parse_PICK_tag(self, node: Node) -> None:
        self._pick_and_parse_children(node)

    # TODO: Refactor _parse_PICK_tag/_parse_SHUFFLE_tag

    def _parse_ROOT_tag(self, node: Node) -> None:
        self._parse_children(node.children)

    def _parse_FREEZE_RANDOM_STATE_tag(self, node: Node) -> None:
        state = random.getstate()
        self._parse_children(node.children)
        random.setstate(state)

    def _parse_TEST_tag(self, node: Node) -> None:
        try:
            if eval(node.arg(0), self.context):
                assert isinstance(node.children[1], Node), repr(node)
                self._parse_children(node.children[1].children)
            else:
                assert isinstance(node.children[2], Node), repr(node)
                self._parse_children(node.children[2].children)
        except Exception:
            print(node.display(color=False))
            raise

    # def _parse_SHARP_tag(self, node: Node) -> None:
    #     """2 sharps (##) will be converted to 1 sharp (#)"""
    #     self.write("#")

    def _parse_ADD_tag(self, node: Node) -> None:
        # a '+' will be displayed at the beginning of the next result if positive ;
        # if the result is negative, nothing will be done, and if null,
        # no result at all will be displayed.
        self.flags.previous_nice_op = NiceOp.ADD

    def _parse_SUB_tag(self, node: Node) -> None:
        # a '-' will be displayed at the beginning of the next result, and the result
        # will be embedded in parentheses if negative.
        self.flags.previous_nice_op = NiceOp.SUB

    def _parse_MUL_tag(self, node: Node) -> None:
        # a '\times' will be displayed at the beginning of the next result, and the result
        # will be embedded in parentheses if negative.
        self.flags.previous_nice_op = NiceOp.MUL

    def _parse_EQUAL_tag(self, node: Node) -> None:
        # Display '=' or '\approx' when a rounded result is requested :
        # if rounded is equal to exact one, '=' is displayed.
        # Else, '\approx' is displayed instead.
        self.flags.previous_nice_op = NiceOp.EQ
        # All other operations (#+, #-, #*) occur just before number, but between `=` and
        # the result, some formatting instructions may occur (like '\fbox{' for example).
        # So, `#=` is used as a temporary marker, and will be replaced by '=' or '\approx' later.
        self.write("#=")

    def _parse_SIGN_tag(self, node: Node) -> None:
        # '>0' or '<0' will be displayed after the next result, depending on it's sign.
        # (If result is zero, this won't do anything.)
        last_value = self.context["_"]
        if last_value > 0:
            self.write(">0")
        elif last_value < 0:
            self.write("<0")

    def _parse_DEBUG_tag(self, node: Optional[Node]) -> None:
        while True:
            msg = "Debug point. Enter command, or quit (q! + ENTER):"
            sep = len(msg) * "="
            print(sep)
            print(msg)
            print(sep)
            command = input("[In]: ")
            print("[Out]: ")
            if command == "q!":
                break
            else:
                try:
                    exec(command, self.context)
                except Exception as e:
                    print("*** ERROR ***")
                    print(e)

    def _parse_PRINT_EVAL_tag(self, node: Node) -> None:
        # print(node.arg(0))
        assert isinstance(node.children[0], Node), repr(node)
        # Backup PTYX_LATEX (i.e. already generated LaTeX code) and reset context.
        assert isinstance(self.context["PTYX_LATEX"], list)
        ptyx_latex_content = self.context["PTYX_LATEX"]
        self.context["PTYX_LATEX"] = []  # Do not use list.clear(), make a new list instead!
        # Interpret code
        self._parse_children(node.children[0].children)
        print("".join(self.context["PTYX_LATEX"]))
        # Restore PTYX_LATEX.
        self.context["PTYX_LATEX"] = ptyx_latex_content

    def _parse_PRINT_tag(self, node: Node) -> None:
        print(node.arg(0))

    def _parse_ALERT_tag(self, node: Node) -> None:
        print(yellow(",".join(pretty_box(node.arg(1), title=node.arg(0)))))

    @staticmethod
    def _exec(code, context):
        """exec is encapsulated in this function to avoid problems
        with free variables in nested functions."""
        exec(code, context)

    @staticmethod
    def _get_python_code_tag(code: str) -> tuple[str, str]:
        tag = ""
        if code.startswith(":"):
            if (i := code.find(":", 1)) != -1:
                tag = code[1:i]
                code = code[i + 1 :]
        return tag, code

    def _exec_python_code(self, code: str, context: dict):
        code = code.replace("\r", "")
        code = code.rstrip().lstrip("\n")
        # Python code snippets can be tagged by using a prefix of the form `:label:`.
        # If an error is raised, the label will be added to the CompilationError instance.
        # This is useful for IDE, to display the faulty piece of code.
        # The IDE can insert a label in each python code snippet.
        # The python code will be transparently executed, yet if an error is raised,
        # the IDE will catch it, read its label to see which code snippet failed,
        # and eventually get the line number and offset to highlight faulty code.
        tag, code = self._get_python_code_tag(code)
        # Indentation test
        initial_indent = len(code) - len(code.lstrip(" "))
        if initial_indent:
            # remove initial indentation
            code = "\n".join(line[initial_indent:] for line in code.split("\n"))
        try:
            self._exec(code, context)
        except Exception as e:  # noqa
            raise PythonBlockError(python_code=code, label=tag, context=self.context) from e
        return code

    def _eval_python_expr(self, code: str) -> object:
        context = self.context
        if not code:
            return ""

        varname = ""
        i = code.find("=")
        if 0 < i < len(code) - 1 and code[i + 1] != "=":
            varname = code[:i].strip()
            if re.match(self.re_varname, varname):
                # This is (probably) a variable name (or a list or dict item reference).
                code = code[i + 1 :]
            else:
                # This is not a variable name.
                varname = ""
        # Last value will be accessible through '_' variable
        if not varname:
            varname = "_"
        if " if " in code and " else " not in code:
            code += " else ''"
        if param["sympy_is_default"]:
            import sympy

            try:
                result = sympy.sympify(code, locals=context)
                if isinstance(result, str):
                    result = result.replace("**", "^")
            except (sympy.SympifyError, AttributeError):
                # sympy.sympify() can't parse attributes and methods inside
                # code for now (AttributeError is raised then).
                result = eval(code, context)
                print("Warning: sympy can't parse %s. " "Switching to standard evaluation mode." % repr(code))
            except Exception:
                # print sorted(context.keys())
                print("Uncatched error when evaluating %s" % repr(code))
                raise
        else:
            result = eval(code, context)
        result = context["_"] = self._apply_flag(result)
        i = varname.find("[")
        # for example, varname == 'mylist[i]' or 'mylist[2]'
        context["LAST"] = result
        if i == -1:
            context[varname] = result
        else:
            key = eval(varname[i + 1 : -1], context)
            varname = varname[:i].strip()
            context[varname][key] = result
        return result

    def _eval_and_format_python_expr(self, code: str) -> str:
        flags = self.flags
        context = self.context
        if not code:
            return ""

        display_result = True
        if code.endswith(";"):
            code = code.rstrip(";")
            display_result = False

        result: object = ""
        for subcode in advanced_split(code, ";", brackets=()):
            result = self._eval_python_expr(subcode)
        # Note that only last result will be displayed.
        # In particular, if code ends with ';', last result will be ''.
        # So, '#{a=5}' and '#{a=5;}' will both affect 5 to `a`,
        # but the second will not display '5' on final document.

        if not display_result:
            return ""

        if flags.is_mul_coeff:
            if result == 1:
                if flags.previous_nice_op == NiceOp.ADD:
                    return "+"
                return ""
            elif result == -1:
                return "-"
            elif result == 0:
                flags.suppress_next_eval = True
                return ""

        if param["sympy_is_default"] and not flags.format_as_str:
            from ptyx.printers import sympy2latex

            latex = sympy2latex(result, flags)
        else:
            latex = str(result)

        def neg(latex_):
            return latex_.lstrip().startswith("-")

        if flags.previous_nice_op == NiceOp.ADD:
            if result == 0:
                latex = ""
            elif not neg(latex):
                latex = "+" + latex
        elif flags.previous_nice_op == NiceOp.MUL:
            latex_list = self.context["PTYX_LATEX"]
            last = self.last_eval_position
            if (
                last.position is not None
                and latex_list[last.position] == last.result
                and all(s.strip() == "" for s in latex_list[last.position + 1 :])
            ):
                if last.result == "1":
                    latex_list[last.position] = ""
                elif last.result == "+1":
                    latex_list[last.position] = "+"
                elif last.result == "-1":
                    latex_list[last.position] = "-"
                elif last.result == "0":
                    return ""
            if neg(latex) or getattr(result, "is_Add", False):
                latex = r"\left(" + latex + r"\right)"
            if not getattr(result, "is_Add", False) and not getattr(result, "is_Symbol", False):
                latex = r"\times " + latex
        elif flags.previous_nice_op == NiceOp.SUB:
            if neg(latex) or getattr(result, "is_Add", False):
                latex = r"\left(" + latex + r"\right)"
            latex = "-" + latex
        elif flags.previous_nice_op == NiceOp.EQ:
            if flags.result_is_exact:
                symb = " = "
            else:
                symb = r" \approx "
            # Search backward for temporary `#=` marker in list, and replace
            # by appropriate symbol.
            for textlist in self.backups + [context["PTYX_LATEX"]]:
                for i, elt in enumerate(reversed(textlist)):
                    if elt == "#=":
                        textlist[len(textlist) - i - 1] = symb
                        return latex
            print("Debug warning: `#=` couldn't be found when scanning context !")
            print("There is most probably a bug in pTyX, entering debugging...")
            self._parse_DEBUG_tag(None)
        return latex

    def _apply_flag(self, result):
        """Apply special parameters like [num], [rand], [floats] or [round] to result.

        Note that [num] and [rand] parameters require that result is iterable,
        otherwise, nothing occurs.
        If result is iterable, an element of result is returned, chosen according
        to current flag."""
        flags = self.flags
        if flags.pick_action != PickItemAction.NONE:
            try:
                if flags.pick_action == PickItemAction.RAND_CHOICE:
                    result = randfunc.randchoice(result)
                elif flags.pick_action == PickItemAction.SELECT_FROM_NUM:
                    result = result[self.NUM % len(result)]
            except (TypeError, IndexError) as e:
                traceback.print_exception(e, limit=2)
                print(
                    "Warning: flags `rand` or `select` have been used for an incorrect data type"
                    " (more details above)."
                )

        if flags.round is not None:
            try:
                if SYMPY_AVAILABLE:
                    round_result = numbers_to_floats(result, ndigits=flags.round)
                else:
                    round_result = round(result, flags.round)
            except ValueError:
                print("** ERROR while rounding value: **")
                print(result)
                print("-----")
                print("".join(self.context["PTYX_LATEX"])[-100:])
                print("-----")
                raise
            if SYMPY_AVAILABLE:
                import sympy
            # noinspection PyUnboundLocalVariable
            if SYMPY_AVAILABLE and isinstance(result, sympy.Basic):
                flags.result_is_exact = {_float_me_if_you_can(elt) for elt in result.atoms()} == {
                    _float_me_if_you_can(elt) for elt in round_result.atoms()
                }
            else:
                flags.result_is_exact = result == round_result
            result = round_result
        elif self.flags.eval_as_float:
            result = numbers_to_floats(result)
        return result


class Compiler:
    """Compiler is the main object of pTyX.

    Usage:

        >>> from ptyx.latex_generator import Compiler
        >>> c = Compiler()
        >>> c.parse('#{a=S(1)/2} + #{b=2} = #{a+b}')
        '\frac{1}{2} + 2 = \frac{5}{2}'

    The following methods are called successively to generate LaTeX code from
    a pTyX file:
        * .read_file(path) will read the content of the file.
        * ._include_subfiles() will parse #INCLUDE{filename} tags and
          include coresponding files content.
        * ._call_extensions() will search for extensions in this content,
          then call the extensions to convert content into plain pTyX code.
        * ._read_seed() will search for a #SEED tag, or give a default seed
          value, used to generate all pseudo-random content later.
        * .generate_syntax_tree() will convert this code into a syntax tree.
          This should be done only once ofr each document, even if multiple
          versions of this document are needed.
        * Finally, .get_latex() will generate and return the LaTeX code.

      Pseudo-random content will depend on the seed (see above), but also
      of the document number, given by `gen.context['PTYX_NUM']`.
      So, changing `gen.context['PTYX_NUM']` enables to generate different
      versions of the same document.
    """

    _state: State

    def __init__(self, *, code=None, path=None):
        self.syntax_tree_generator = SyntaxTreeGenerator()
        self.latex_generator = LatexGenerator(self)
        self.reset()
        if path is not None or code is not None:
            self.load(path=path, code=code)
            self.preparse()
            self.generate_syntax_tree()

    def reset(self) -> None:
        self._state = {}
        # Make SyntaxTreeGenerator context free ?
        self.syntax_tree_generator.reset()
        self.latex_generator.reset()

    def read_code(self, code: str) -> None:
        """Feed compiler with given code."""
        self._state["path"] = None
        self._state["input"] = code

    def read_file(self, path: Union[Path, str]) -> None:
        """Feed compiler with given file code."""
        # We must use `os.path.abspath()` here:
        #   - `Path.absolute()` would not normalize the path,
        #      i.e. ".." and "." are not interpreted,
        #      obsfucating any bug report.
        #   - `Path.resolve()` or `os.path.realpath()` would
        #      resolve symlinks, and so change the folder where
        #      all the files are generated as a side effect.
        path = Path(os.path.abspath(os.path.expanduser(path)))
        self._state["path"] = path
        with open(path, "r") as input_file:
            self._state["input"] = input_file.read()

    @property
    def dir_path(self) -> Path:
        """Return input ptyx file directory, if any, or current working directory else."""
        file_path = self.file_path
        return Path.cwd() if file_path is None else file_path.parent

    def _resolve_input_file_path(self, path: Union[str, Path]) -> Path:
        """Interpret `path` relatively to input ptyx file."""
        if isinstance(path, str):
            path = Path(path.strip())
        path = path.expanduser()  # do NOT resolve yet !
        if not path.is_absolute():
            path = self.dir_path / path
        return path

    def _include_subfiles(self, code: str) -> str:
        """Parse all #INCLUDE tags, then include corresponding files content."""

        def include(match: re.Match) -> str:
            path = self._resolve_input_file_path(match.group(1))
            with open(path) as file:
                return f"\n#APART#INCLUDE_START{{{path}}}{{{match.start()}}}\n{file.read()}#END_APART#INCLUDE_END\n"

        # noinspection RegExpRedundantEscape
        return re.sub(r"#INCLUDE\{([^}]+)\}", include, code)

    def _call_extensions(self, code: str) -> tuple[str, dict[str, ModuleType]]:
        """Search for extensions (#LOAD{name} tags), then call them."""
        # First, we search if some extensions must be load.
        # This must be done at the very beginning, since extensions may
        # define their own specialized language, to be converted to
        # valid pTyX code (and then to LaTeX).

        # Search for #LOAD{} tags and list corresponding extensions.
        # noinspection RegExpRedundantEscape
        extensions_list = re.findall(r"#LOAD\{\s*(\w+)\s*\}", code)

        extensions: Dict[Tag, ModuleType] = {}
        tags_syntax: Dict[Tag, TagSyntax] = {}
        tags_source: Dict[Tag, str] = {}
        latex_generator_extensions = []
        for extension_name in extensions_list:
            print(f"Loading extension '{extension_name}'...")
            try:
                extensions[extension_name] = import_module(f"ptyx.extensions.{extension_name}")
            except ImportError:
                # Try to find a matching registered plugin.
                for entry_point in metadata.entry_points(group="ptyx.extensions"):
                    if entry_point.name == extension_name:
                        extensions[extension_name] = import_module(entry_point.value)
                        break
                else:
                    traceback.print_exc()
                    raise PtyxExtensionNotFound(f"Extension {extension_name} not found.")
            try:
                # This extension may define a function `extend_compiler()` to customize the compiler.
                extensions_dict: CompilerExtension = getattr(extensions[extension_name], "extend_compiler")()
                # This extension may subclass LatexGenerator class, notably to handle new tags.
                if "latex_generator" in extensions_dict:
                    latex_generator_extensions.append(extensions_dict["latex_generator"])
                # Test if this extension defines new tags.
                ext_tags = extensions_dict.get("tags", {})
                for tag, syntax in ext_tags.items():
                    # Test for conflict between extensions.
                    if tag in tags_source:
                        raise NameError(
                            f"Extension {extension_name} tries to define tag {tag}, "
                            f"which was already defined by extension {tags_source[tag]}."
                        )
                    # Tag not already declared, everything seems OK.
                    tags_source[tag] = extension_name
                    tags_syntax[tag] = syntax
            except AttributeError:
                pass
        # Update LatexGenerator.
        if latex_generator_extensions:
            # TODO: Test for conflicting methods ?
            class CustomLatexGenerator(*reversed(latex_generator_extensions)):  # type: ignore
                """Custom Latex Generator dynamically generated using loaded extensions."""

            self.latex_generator = CustomLatexGenerator(self)
        # Load new tags. This must be done *AFTER* the redefinition of self.latex_generator,
        # so as to update the parser of the *new* latex generator, and not the old one.
        self.add_new_tags(*tags_syntax.items())
        for name in extensions_list:
            # execute `main()` function of extension.
            if hasattr(extensions[name], "main"):
                code = extensions[name].main(code, self)
        return code, extensions

    def _read_seed(self, code: str) -> Tuple[str, Optional[int]]:
        """Extract seed value from code, searching for #SEED{num} tag.

        Return the code without the #SEED{...} tag, and the seed value (if any, `None` else).
        """
        counter = 0
        value = None

        def seed(match):
            nonlocal counter, value
            value = int(match.group(1))
            counter += 1
            return ""

        # noinspection RegExpRedundantEscape
        code = re.sub(r"#SEED\{\s*(\d+)\s*\}", seed, code)
        if counter == 0:
            path = self._state.get("path")
            value = zlib.adler32(str(path).encode("utf8"))
            print(f"Warning: #SEED not found, using ptyx file path '{path}' to generate seed: {value}.")
        elif counter > 1:
            print(f"Warning: multiple #SEED found, only last one will be used: {value}.")
        return code, value

    @staticmethod
    def _assert_no_comment(code: str) -> None:
        """AssertionError raised if comments remain."""
        comments = any((" # " in code, "\n# " in code, code.startswith("# ")))
        if comments:
            print(code)
        assert not comments, "There should be no remaining comment. Maybe a problem with an extension ?"

    def preparse(self) -> None:
        code = self._state.get("input")
        if code is None:
            raise RuntimeError("Compiler.read_code() or Compiler.read_file() must be run first.")
        remove_comments = self.syntax_tree_generator.remove_comments
        # TODO: do not remove comments before resolving imports, since this prevent tracebacks from providing
        #  accurate positions.
        #  Instead, self._include_subfiles() should be improved to ignore imports in lines starting with a #.
        # Remove comments first, so one can comment a file inclusion for example.
        code = remove_comments(code)
        assert isinstance(code, str)
        self._assert_no_comment(code)
        code = remove_comments(self._include_subfiles(code))
        self._state["after_include"] = code
        assert isinstance(code, str)
        code, extensions = self._call_extensions(code)
        code, seed = self._read_seed(code)
        # Remove any comment that may have been generated by an extension.
        code = remove_comments(code)
        assert isinstance(code, str)
        self._state["plain_ptyx_code"] = code
        self._state["loaded_extensions"] = extensions
        # Set the seed used for pseudo-random numbers generation.
        # (The seed value is set in the ptyx file using special tag #SEED{}).
        self._state["seed"] = seed
        self._assert_no_comment(code)
        assert "#INCLUDE{" not in code
        # assert "#LOAD{" not in code
        assert "#SEED{" not in code
        # Save pTyX code generated by extensions (this is used for debugging,
        # but if needed extensions can also save some data this way using #COMMENT tag).
        # If input file was /path/to/file/myfile.ptyx,
        # plain pTyX code is saved in /path/to/file/.myfile.ptyx.plain-ptyx
        if extensions:
            path = self._state.get("path")
            if path is not None:
                filename = join(dirname(path), ".%s.plain-ptyx" % basename(path))
                with open(filename, "w") as f:
                    f.write(code)

    def generate_syntax_tree(self) -> None:
        """Generate the syntax tree."""
        code = self._state.get("plain_ptyx_code")
        if code is None:
            raise RuntimeError("Compiler.preparse() must be run first.")
        self._state["syntax_tree"] = self.syntax_tree_generator.generate_tree(code)

    def get_latex(self, **context) -> str:
        """Compile pTyX code and return LaTeX code.

        `Compiler.generate_syntax_tree()` must be run first.
        """
        tree = self._state.get("syntax_tree")
        if tree is None:
            raise RuntimeError("`Compiler.generate_syntax_tree()` must be run first.")
        gen = self.latex_generator
        gen.clear(main_file=self._state["path"])
        gen.context.update(context)
        seed = self._state["seed"] + gen.NUM
        randfunc.set_seed(seed)
        try:
            gen.parse_node(tree)
        except Exception as e:
            print("\n*** Error occurred while generating LaTeX code. ***")
            print("This was last generated LaTeX code for debugging purpose:")
            print(80 * "-")
            print("... " + "".join(gen.context["PTYX_LATEX"][-10:]))
            print(80 * "-")
            print("")
            if hasattr(e, "pretty_report"):
                print(e.pretty_report)
            raise
        latex = gen.read()
        # if "API_VERSION" not in gen.context:
        #     print("Warning: no API version specified. This may be an old pTyX file.")
        return latex

    def add_new_tags(self, *tags: tuple[str, tuple[int, int, list[str] | None]]) -> None:
        """Add ability for extensions to extend syntax, adding new tags."""
        for name, syntax in tags:
            self.syntax_tree_generator.tags[name] = syntax
            self.latex_generator.parser.tags[name] = syntax
        if tags:
            self.syntax_tree_generator.update_tags()
            self.latex_generator.parser.update_tags()

    def load(self, *, code: str = None, path: Union[Path, str] = None) -> None:
        """Reset state, parse pTyX code and generate the syntax tree.

        One may provide either directly the pTyX code, or the path of a pTyX file to be read.

        If both `code` and `path` are provided, the compiler acts as if `code` was the content
        of the pTyX file located at `path`."""
        self.reset()
        if path is None and code is None:
            raise TypeError(
                "You must provide either the pTyX code or the path to a pTyX file:\n"
                'compiler.load(code="...") or compiler.load(path="/path/to/file.ptyx")'
            )
        if path is not None:
            self.read_file(path)
        if code is not None:
            self.read_code(code)
        self.preparse()
        self.generate_syntax_tree()

    def parse(self, *, code: str = None, path: Union[Path, str] = None, **context) -> str:
        """Convert ptyx code to plain LaTeX in one shot.

        This is mainly used for testing (in unit tests or in interactive mode).

        One may provide either directly the pTyX code, or the path of a pTyX file to be read.

        If both `code` and `path` are provided, the compiler acts as if `code` was the content
        of the pTyX file located at `path`.
        """
        self.load(code=code, path=path)
        latex = self.get_latex(**context)
        return latex

    @property
    def syntax_tree(self) -> Node:
        return self._state["syntax_tree"]

    @property
    def seed(self) -> Optional[int]:
        return self._state["seed"]

    @property
    def file_path(self) -> Optional[Path]:
        return self._state["path"]

    @property
    def plain_ptyx_code(self) -> Optional[str]:
        return self._state["plain_ptyx_code"]

    @property
    def loaded_extensions(self) -> List[str]:
        return list(self._state["loaded_extensions"].keys())
