import re
import sys
import warnings
from collections.abc import Callable, Generator
from itertools import zip_longest
from typing import TYPE_CHECKING, Any, NoReturn, cast

from ._key_path import NodeInterpolationKey, split_key
from ._missing import _decode_missing_escape, _EscapedMissing, _has_missing_spelling
from .errors import InterpolationResolutionError
from .vendor.antlr4 import TerminalNode  # type: ignore[attr-defined]

if TYPE_CHECKING:
    from .base import Node  # noqa F401

try:
    from omegaconf.grammar.gen.OmegaConfGrammarLexer import OmegaConfGrammarLexer
    from omegaconf.grammar.gen.OmegaConfGrammarParser import OmegaConfGrammarParser
    from omegaconf.grammar.gen.OmegaConfGrammarParserVisitor import (
        OmegaConfGrammarParserVisitor,
    )

except ModuleNotFoundError:  # pragma: no cover
    print(
        "Error importing OmegaConf's generated parsers, run `python setup.py antlr` to regenerate.",
        file=sys.stderr,
    )
    sys.exit(1)


class GrammarVisitor(OmegaConfGrammarParserVisitor):
    node_interpolation_callback: (
        Callable[[NodeInterpolationKey, set[int] | None], "Node | None"] | None
    )
    resolver_interpolation_callback: Callable[..., Any] | None
    memo: set[int] | None

    def __init__(
        self,
        node_interpolation_callback: (
            Callable[
                [NodeInterpolationKey, set[int] | None],
                "Node | None",
            ]
            | None
        ),
        resolver_interpolation_callback: Callable[..., Any] | None,
        memo: set[int] | None,
        **kw: dict[Any, Any],
    ):
        """
        Constructor.

        :param node_interpolation_callback: Callback function that is called when
            needing to resolve a node interpolation. Its key includes the raw path
            and the parsed segments (ex: `("foo", "bar")`).

        :param resolver_interpolation_callback: Callback function that is called when
            needing to resolve a resolver interpolation. This function should accept
            three keyword arguments: `name` (str, the name of the resolver),
            `args` (tuple, the inputs to the resolver), and `args_str` (tuple,
            the string representation of the inputs to the resolver).

        :param kw: Additional keyword arguments to be forwarded to parent class.
        """
        super().__init__(**kw)
        self.node_interpolation_callback = node_interpolation_callback
        self.resolver_interpolation_callback = resolver_interpolation_callback
        self.memo = memo

    def aggregateResult(self, aggregate: list[Any], nextResult: Any) -> list[Any]:
        raise NotImplementedError

    def defaultResult(self) -> NoReturn:
        # Raising an exception because not currently used (like `aggregateResult()`).
        raise NotImplementedError

    def visitConfigKey(self, ctx: OmegaConfGrammarParser.ConfigKeyContext) -> str:
        from ._utils import _get_value

        # interpolation | ID | INTER_KEY
        assert ctx.getChildCount() == 1
        child = ctx.getChild(0)
        if isinstance(child, OmegaConfGrammarParser.InterpolationContext):
            child_text = child.getText()
            res = _get_value(self.visitInterpolation(child))
            if type(res) is int:
                return str(res)
            elif not isinstance(res, str):
                raise InterpolationResolutionError(
                    f"The following interpolation is used to denote a config key and "
                    f"thus should return a string or integer, but instead returned `{res}` of "
                    f"type `{type(res)}`: {child_text}"
                )
            return res
        else:
            assert isinstance(child, TerminalNode) and isinstance(
                child.symbol.text,  # type: ignore[attr-defined]
                str,
            )
            text = child.symbol.text  # type: ignore[attr-defined]
            return re.sub(r"\\([\\.\[\]:=])", r"\1", text)

    def visitConfigValue(self, ctx: OmegaConfGrammarParser.ConfigValueContext) -> Any:
        # text EOF
        assert ctx.getChildCount() == 2
        return self.visit(ctx.getChild(0))

    def visitDictKey(self, ctx: OmegaConfGrammarParser.DictKeyContext) -> Any:
        return self._createPrimitive(ctx)

    def visitDictContainer(
        self, ctx: OmegaConfGrammarParser.DictContainerContext
    ) -> dict[Any, Any]:
        # BRACE_OPEN (dictKeyValuePair (COMMA dictKeyValuePair)*)? BRACE_CLOSE
        assert ctx.getChildCount() >= 2
        return dict(
            self.visitDictKeyValuePair(
                cast(OmegaConfGrammarParser.DictKeyValuePairContext, ctx.getChild(i))
            )
            for i in range(1, ctx.getChildCount() - 1, 2)
        )

    def visitElement(self, ctx: OmegaConfGrammarParser.ElementContext) -> Any:
        # primitive | quotedValue | listContainer | dictContainer
        assert ctx.getChildCount() == 1
        value = self.visit(ctx.getChild(0))
        spelling = ctx.getText()
        if len(spelling) >= 2 and spelling[0] in "\"'" and spelling[-1] == spelling[0]:
            spelling = spelling[1:-1]
        if isinstance(value, str) and _has_missing_spelling(spelling):
            return _decode_missing_escape(spelling)
        return value

    def visitInterpolation(
        self, ctx: OmegaConfGrammarParser.InterpolationContext
    ) -> Any:
        assert ctx.getChildCount() == 1  # interpolationNode | interpolationResolver
        return self.visit(ctx.getChild(0))

    def visitInterpolationNode(
        self, ctx: OmegaConfGrammarParser.InterpolationNodeContext
    ) -> "Node | None":
        # INTER_OPEN
        # DOT*                                                     // relative interpolation?
        # (configKey | BRACKET_OPEN configKey BRACKET_CLOSE)       // foo, [foo]
        # (DOT configKey | BRACKET_OPEN configKey BRACKET_CLOSE)*  // .foo, [foo], .foo[bar], [foo].bar[baz]
        # INTER_CLOSE;

        assert ctx.getChildCount() >= 3

        inter_key_tokens = []  # original spelling for relative-root selection and errors
        parts: list[str] = []
        relative_dots = 0
        first_key = True
        in_brackets = False
        for child in ctx.getChildren():
            if isinstance(child, TerminalNode):
                s = child.symbol  # type: ignore[attr-defined]
                if s.type in [
                    OmegaConfGrammarLexer.DOT,
                    OmegaConfGrammarLexer.BRACKET_OPEN,
                    OmegaConfGrammarLexer.BRACKET_CLOSE,
                ]:
                    inter_key_tokens.append(s.text)
                    if first_key and s.type == OmegaConfGrammarLexer.DOT:
                        relative_dots += 1
                    if s.type == OmegaConfGrammarLexer.BRACKET_OPEN:
                        in_brackets = True
                    elif s.type == OmegaConfGrammarLexer.BRACKET_CLOSE:
                        in_brackets = False
                else:
                    assert s.type in (
                        OmegaConfGrammarLexer.INTER_OPEN,
                        OmegaConfGrammarLexer.INTER_CLOSE,
                    )
            else:
                assert isinstance(child, OmegaConfGrammarParser.ConfigKeyContext)
                value = self.visitConfigKey(child)
                if isinstance(
                    child.getChild(0), OmegaConfGrammarParser.InterpolationContext
                ):
                    inter_key_tokens.append(value)
                    expanded = (
                        split_key(f"[{value}]") if in_brackets else split_key(value)
                    )
                    if first_key and not in_brackets:
                        leading_dots = len(value) - len(value.lstrip("."))
                        relative_dots += leading_dots
                        expanded = (
                            []
                            if leading_dots == len(value)
                            else expanded[leading_dots:]
                        )
                    parts.extend(expanded)
                else:
                    inter_key_tokens.append(child.getText())
                    parts.append(value)
                first_key = not parts

        inter_key = NodeInterpolationKey(
            "".join(inter_key_tokens), tuple(parts), relative_dots
        )
        assert self.node_interpolation_callback is not None
        return self.node_interpolation_callback(inter_key, self.memo)

    def visitInterpolationResolver(
        self, ctx: OmegaConfGrammarParser.InterpolationResolverContext
    ) -> Any:
        # INTER_OPEN resolverName COLON sequence? BRACE_CLOSE
        assert 4 <= ctx.getChildCount() <= 5

        resolver_name = self.visit(ctx.getChild(1))
        maybe_seq = ctx.getChild(3)
        args = []
        args_str = []
        if isinstance(maybe_seq, TerminalNode):  # means there are no args
            assert maybe_seq.symbol.type == OmegaConfGrammarLexer.BRACE_CLOSE  # type: ignore[attr-defined]
        else:
            assert isinstance(maybe_seq, OmegaConfGrammarParser.SequenceContext)
            for val, txt in self.visitSequence(maybe_seq):
                args.append(val)
                args_str.append(txt)

        assert self.resolver_interpolation_callback is not None
        return self.resolver_interpolation_callback(
            name=resolver_name,
            args=tuple(args),
            args_str=tuple(args_str),
        )

    def visitDictKeyValuePair(
        self, ctx: OmegaConfGrammarParser.DictKeyValuePairContext
    ) -> tuple[Any, Any]:
        from ._utils import _get_value

        assert ctx.getChildCount() == 3  # dictKey COLON element
        key = self.visit(ctx.getChild(0))
        colon = ctx.getChild(1)
        assert (
            isinstance(colon, TerminalNode)
            and colon.symbol.type == OmegaConfGrammarLexer.COLON  # type: ignore[attr-defined]
        )
        child = ctx.getChild(2)
        assert child is not None  # help type checkers
        value = _get_value(self.visitElement(child))
        return key, value

    def visitListContainer(
        self, ctx: OmegaConfGrammarParser.ListContainerContext
    ) -> list[Any]:
        # BRACKET_OPEN sequence? BRACKET_CLOSE;
        assert ctx.getChildCount() in (2, 3)
        if ctx.getChildCount() == 2:
            return []
        sequence = ctx.getChild(1)
        assert isinstance(sequence, OmegaConfGrammarParser.SequenceContext)
        return list(val for val, _ in self.visitSequence(sequence))  # ignore raw text

    def visitPrimitive(self, ctx: OmegaConfGrammarParser.PrimitiveContext) -> Any:
        return self._createPrimitive(ctx)

    def visitQuotedValue(self, ctx: OmegaConfGrammarParser.QuotedValueContext) -> str:
        # (QUOTE_OPEN_SINGLE | QUOTE_OPEN_DOUBLE) text? MATCHING_QUOTE_CLOSE
        n = ctx.getChildCount()
        assert n in [2, 3]
        if n == 2:
            return ""
        child = ctx.getChild(1)
        assert isinstance(child, OmegaConfGrammarParser.TextContext)
        return self._unescape(list(child.getChildren()))

    def visitResolverName(self, ctx: OmegaConfGrammarParser.ResolverNameContext) -> str:
        from ._utils import _get_value

        # (interpolation | ID) (DOT (interpolation | ID))*
        assert ctx.getChildCount() >= 1
        items = []
        for child in list(ctx.getChildren())[::2]:
            if isinstance(child, TerminalNode):
                assert child.symbol.type == OmegaConfGrammarLexer.ID  # type: ignore[attr-defined]
                items.append(child.symbol.text)  # type: ignore[attr-defined]
            else:
                assert isinstance(child, OmegaConfGrammarParser.InterpolationContext)
                item = _get_value(self.visitInterpolation(child))
                if not isinstance(item, str):
                    raise InterpolationResolutionError(
                        f"The name of a resolver must be a string, but the interpolation "
                        f"{child.getText()} resolved to `{item}` which is of type "
                        f"{type(item)}"
                    )
                items.append(item)
        return ".".join(items)

    def visitSequence(
        self, ctx: OmegaConfGrammarParser.SequenceContext
    ) -> Generator[Any, None, None]:
        from ._utils import _get_value

        # (element (COMMA element?)*) | (COMMA element?)+
        assert ctx.getChildCount() >= 1

        # DEPRECATED: remove in 2.2 (revert #571)
        def empty_str_warning() -> None:
            txt = ctx.getText()
            warnings.warn(
                f"In the sequence `{txt}` some elements are missing: please replace "
                f"them with empty quoted strings. "
                f"See https://github.com/hydra-ecosystem/omegaconf/issues/572 for details.",  # noqa: E231
                category=UserWarning,
            )

        is_previous_comma = True  # whether previous child was a comma (init to True)
        for child in ctx.getChildren():
            if isinstance(child, OmegaConfGrammarParser.ElementContext):
                # Also preserve the original text representation of `child` so
                # as to allow backward compatibility with old resolvers (registered
                # with `legacy_register_resolver()`). Note that we cannot just cast
                # the value to string later as for instance `null` would become "None".
                yield _get_value(self.visitElement(child)), child.getText()
                is_previous_comma = False
            else:
                assert (
                    isinstance(child, TerminalNode)
                    and child.symbol.type == OmegaConfGrammarLexer.COMMA  # type: ignore[attr-defined]
                )
                if is_previous_comma:
                    empty_str_warning()
                    yield "", ""
                else:
                    is_previous_comma = True
        if is_previous_comma:
            # Trailing comma.
            empty_str_warning()
            yield "", ""

    def visitSingleElement(
        self, ctx: OmegaConfGrammarParser.SingleElementContext
    ) -> Any:
        # element EOF
        assert ctx.getChildCount() == 2
        return self.visit(ctx.getChild(0))

    def visitText(self, ctx: OmegaConfGrammarParser.TextContext) -> Any:
        # (interpolation | ANY_STR | ESC | ESC_INTER | TOP_ESC | QUOTED_ESC)+

        # Single interpolation? If yes, return its resolved value "as is".
        if ctx.getChildCount() == 1:
            c = ctx.getChild(0)
            if isinstance(c, OmegaConfGrammarParser.InterpolationContext):
                return self.visitInterpolation(c)

        # Otherwise, concatenate string representations together.
        return self._unescape(list(ctx.getChildren()))

    def _createPrimitive(
        self,
        ctx: (
            OmegaConfGrammarParser.PrimitiveContext
            | OmegaConfGrammarParser.DictKeyContext
        ),
    ) -> Any:
        # (ID | NULL | INT | FLOAT | BOOL | UNQUOTED_CHAR | COLON | ESC | WS | interpolation)+
        if ctx.getChildCount() == 1:
            child = ctx.getChild(0)
            if isinstance(child, OmegaConfGrammarParser.InterpolationContext):
                return self.visitInterpolation(child)
            assert isinstance(child, TerminalNode)
            symbol = child.symbol  # type: ignore[attr-defined]
            # Parse primitive types.
            if symbol.type in (
                OmegaConfGrammarLexer.ID,
                OmegaConfGrammarLexer.UNQUOTED_CHAR,
                OmegaConfGrammarLexer.COLON,
            ):
                return symbol.text
            elif symbol.type == OmegaConfGrammarLexer.NULL:
                return None
            elif symbol.type == OmegaConfGrammarLexer.INT:
                return int(symbol.text)
            elif symbol.type == OmegaConfGrammarLexer.FLOAT:
                return float(symbol.text)
            elif symbol.type == OmegaConfGrammarLexer.BOOL:
                return symbol.text.lower() == "true"
            elif symbol.type == OmegaConfGrammarLexer.ESC:
                return self._unescape([child])
            elif symbol.type == OmegaConfGrammarLexer.WS:  # pragma: no cover
                # A single WS should have been "consumed" by another token.
                raise AssertionError("WS should never be reached")
            assert False, symbol.type
        # Concatenation of multiple items ==> un-escape the concatenation.
        return self._unescape(list(ctx.getChildren()))

    def _unescape(
        self,
        seq: list[TerminalNode | OmegaConfGrammarParser.InterpolationContext],
    ) -> str:
        """
        Concatenate all symbols / interpolations in `seq`, unescaping symbols as needed.

        Interpolations are resolved and cast to string *WITHOUT* escaping their result
        (it is assumed that whatever escaping is required was already handled during the
        resolving of the interpolation).
        """
        from ._utils import _get_value

        chrs = []
        has_escaped_missing = False
        for node, next_node in zip_longest(seq, seq[1:]):
            if isinstance(node, TerminalNode):
                s = node.symbol  # type: ignore
                if s.type == OmegaConfGrammarLexer.ESC_INTER:
                    # `ESC_INTER` is of the form `\\...\${`: the formula below computes
                    # the number of characters to keep at the end of the string to remove
                    # the correct number of backslashes.
                    text = s.text[-(len(s.text) // 2 + 1) :]
                elif (
                    # Character sequence identified as requiring un-escaping.
                    s.type == OmegaConfGrammarLexer.ESC
                    or (
                        # At top level, we need to un-escape backslashes that precede
                        # an interpolation.
                        s.type == OmegaConfGrammarLexer.TOP_ESC
                        and isinstance(
                            next_node, OmegaConfGrammarParser.InterpolationContext
                        )
                    )
                    or (
                        # In a quoted sring, we need to un-escape backslashes that
                        # either end the string, or are followed by an interpolation.
                        s.type == OmegaConfGrammarLexer.QUOTED_ESC
                        and (
                            next_node is None
                            or isinstance(
                                next_node, OmegaConfGrammarParser.InterpolationContext
                            )
                        )
                    )
                ):
                    text = s.text[1::2]  # un-escape the sequence
                else:
                    text = s.text  # keep the original text
            else:
                assert isinstance(node, OmegaConfGrammarParser.InterpolationContext)
                value = _get_value(self.visitInterpolation(node))
                has_escaped_missing |= isinstance(value, _EscapedMissing)
                text = str(value)
            chrs.append(text)

        result = "".join(chrs)
        if has_escaped_missing and _has_missing_spelling(result):
            return _EscapedMissing(result)
        return result