"""
    Support for traversing the IR, often to search for information.
"""
from typing import Callable, Dict, Optional, Type

from ...metamodel import metamodel as mm
from ....util.structures import OrderedSet
from ...metamodel.rewriter import _ATOMS, NO_WALK, WALK_FIELDS

#--------------------------------------------------
# Visitor
#--------------------------------------------------
# Similar to Walker, but with enter/exit hooks.

class Visitor():
    visits = 0
    _visitors: Dict[Type[mm.Node], Callable] = {}
    __slots__ = ("_seen_ids",)

    def __init__(self):
        # ids of nodes on the current recursion path (for cycle detection)
        self._seen_ids: set[int] = set()

    def __call__(self, obj):
        # _seen_ids is properly maintained by add/remove in _visit_node
        try:
            self._visit(obj,)
        except StopIteration:
            pass
        return obj

    def enter(self, node: mm.Node, parent: Optional[mm.Node]=None) -> "Visitor":
        return self

    def exit(self, node: mm.Node, parent: Optional[mm.Node]=None) -> mm.Node:
        return node

    def _visit(self, x, parent: Optional[mm.Node]=None):
        self.enter(x, parent)
        if isinstance(x, mm.Node):
            node = self._visit_node(x, parent)
            return self.exit(node, parent)
        if isinstance(x, _ATOMS):
            return x
        if isinstance(x, (tuple, list, OrderedSet)):
            for v in x:
                self._visit(v, parent)
            return x
        if isinstance(x, dict):
            for v in x.values():
                self._visit(v, parent)
            return x
        return x

    def _visit_node(self, node: mm.Node, parent: Optional[mm.Node]=None):
        nid = id(node)
        if nid in self._seen_ids:  # cycle guard by identity
            return node

        Visitor.visits += 1
        self._seen_ids.add(nid)

        cls = type(node)
        self._visitors[cls](self, node, parent)

        if nid in self._seen_ids:
            self._seen_ids.remove(nid)
        return node

    #------------------------------------------------------
    # Generate visitors
    #------------------------------------------------------

    @classmethod
    def generate_visitors(cls):
        dispatch: Dict[Type[mm.Node], Callable] = {}

        for node in WALK_FIELDS.keys():
            base_name = node.__name__.lower()
            enter_name = f"enter_{base_name}"
            visit_name = f"visit_{base_name}"
            exit_name = f"exit_{base_name}"

            # default no-op enter handler
            def _no_op(self, node, parent=None):
                return None
            setattr(cls, enter_name, _no_op)
            setattr(cls, visit_name, _no_op)
            setattr(cls, exit_name, _no_op)

            # Generate the per-class walker via exec for speed
            ns = {"NO_WALK": NO_WALK}
            lines =      ["def visitor(self, node, parent=None):"]
            lines.append(f"    no_walk = self.{enter_name}(node, parent)")
            lines.append("    if no_walk is not NO_WALK:")
            lines.extend(f"        self._visit(node.{fld}, node)" for fld in WALK_FIELDS.get(node, ()))
            lines.append(f"        self.{visit_name}(node, parent)")
            lines.append(f"        self.{exit_name}(node, parent)")
            exec("\n".join(lines), ns)
            dispatch[node] = ns["visitor"] # type: ignore

        cls._visitors = dispatch

Visitor.generate_visitors()
