from __future__ import annotations

from typing import Any, Final, Optional

from jinja2 import Environment
from jinja2.nodes import Call, Const, For, Getattr, Getitem, Name, Node, Tuple
from jinja2.visitor import NodeVisitor
from typing_extensions import TypeAlias

VariablePath: TypeAlias = tuple[str, ...]  # e.g. (ntp_servers, ip) implies ntp_servers.ip
IterablePath: TypeAlias = tuple[str, ...]  # e.g. (users,params) implies a nested loop for users: for params: x
AliasMap: TypeAlias = dict[str, list[str]]  # e.g. loop var name -> iterable path, e.g. {"param": ["users", "params"]}

DYNAMIC_KEY_PLACEHOLDER: Final[str] = "<key>"  # placeholder for non-constant/unknown index keys


# noinspection PyPep8Naming
class VariableCollector(NodeVisitor):
    """
    Collects all variable paths and iterable roots from a Jinja2 AST.
    - Uses a stack to detect when a `Name` node is part of a chain.
    - Tracks loop aliases to resolve variable paths within loops
    - Returns sets of absolute paths for both variables and iterables.
    """

    def __init__(self) -> None:
        self.paths: set[VariablePath] = set()
        self.iterables: set[IterablePath] = set()
        self._node_stack: Final[list[Node]] = []
        self._alias_stack: list[AliasMap] = []

    def visit(self, node: Node, *args: Any, **kwargs: Any) -> None:
        """
        Visit a node while tracking the parent on a simple stack.
        """
        self._node_stack.append(node)
        super().visit(node)
        self._node_stack.pop()

    def visit_Name(self, node: Name) -> None:
        """
        Visit each Name node, which is used in Jinja to lookup or store a value in a name.
        Any flat variable like "{{ abc }}" will be represented as a Name node.
        """

        # Ignore any non-load contexts, e.g. function params or assignments.
        if node.ctx != "load":
            return

        # Ignore if part of a chained access, e.g. abc.def or abc[def], otherwise this would double-count.
        parent = self._get_node_parent()
        if isinstance(parent, (Getattr, Getitem)):
            return

        # Resolve alias (if any) and record the absolute path
        path = self._resolve_alias([node.name])
        self.paths.add(tuple(path))

    def visit_Getattr(self, node: Getattr) -> None:
        """
        Visit each Getattr node, which represents attribute access in Jinja.
        There is some special case handling to avoid recording paths for calls to items/keys/values.
        E.g. "{{ foo.bar }}" is represented as Getattr(node=Name(foo), attr=bar).
        """
        # Detect if this is a call to items/keys/values
        parent = self._get_node_parent()
        is_callee = isinstance(parent, Call) and parent.node is node and node.attr in {"items", "keys", "values"}

        # Skip recording this path if it's the callee of an items/keys/values call.
        if not is_callee:
            path = self._expr_to_path(node)
            if path is not None:
                self.paths.add(tuple(path))

        # Call default visitor which traverses child nodes.
        self.generic_visit(node)

    def visit_Getitem(self, node: Getitem) -> None:
        """
        Visit each Getitem node, which represents index access in Jinja.
        E.g. "{{ foo['bar'] }}" is represented as Getitem(node=Name(foo), arg=Const('bar')).
        """
        path = self._expr_to_path(node)
        if path is not None:
            self.paths.add(tuple(path))

        # Call default visitor which traverses child nodes.
        self.generic_visit(node)

    def visit_For(self, node: For) -> None:
        """
        Visit each For node, which represents a for-loop in Jinja, e.g. "{% for u in users %}...{% endfor %}".
        1. Collect the iterable path (e.g. "users").
        2. Push an alias mapping for the loop target(s) to the iterable path.
        3. Visit the loop body with the alias in scope.
        4. Visit the else body if present.
        5. Pop the alias mapping.
        """
        # Determine the iterable path, e.g. "users" in "for u in users".
        iterable_path = self._iterable_path_from_iter(node.iter)
        if iterable_path:
            self.iterables.add(tuple(iterable_path))
            self.paths.add(tuple(iterable_path))

        # Visit the iterable expression first to collect any nested variables.
        self.visit(node.iter)

        # Determine loop target names, e.g. "u" in "for u in users".
        targets = self._target_names(node.target)

        # Push a temporary alias scope for the loop targets before visiting the body.
        pushed = False
        if targets and iterable_path is not None:
            # Use special handling for dict-based loops with .items().
            # The key alias is being mapped to iterable/<key-alias> for clarity.
            # The remaining aliases (= the object values) map to the iterable path directly.
            alias_scope: AliasMap = {}
            if self._is_items_call(node.iter) and len(targets) >= 2:
                alias_scope = {
                    targets[0]: list(iterable_path) + [targets[0]],
                    **{target: list(iterable_path) for target in targets[1:]},
                }
            else:
                alias_scope = {target: list(iterable_path) for target in targets}

            self._alias_stack.append(alias_scope)
            pushed = True

        # Visit the loop body, and then the else body if present.
        for child in node.body:
            self.visit(child)
        else_body = getattr(node, "else_", []) or []
        for child in else_body:
            self.visit(child)

        # Remove the alias scope after visiting the loop.
        if pushed:
            self._alias_stack.pop()

    def _is_alias_name(self, name: str) -> bool:
        """
        Returns true if the name is currently bound as a loop alias.
        E.g. in "{% for u in users %} {{ u.name }} {% endfor %}", "u" is an alias for "users".
        """
        for scope in reversed(self._alias_stack):
            if name in scope:
                return True

        return False

    def _resolve_alias(self, path: list[str]) -> list[str]:
        """
        Resolves the first element of the path if it is a loop alias.
        E.g. in "{% for u in users %} {{ u.name }} {% endfor %}", "u.name" resolves to "users.name".
        """
        if not path:
            return path

        head = path[0]
        for scope in reversed(self._alias_stack):
            if head in scope:
                return list(scope[head]) + path[1:]

        return path

    def _get_node_parent(self) -> Optional[Node]:
        """
        Get the parent node of the current node on the stack, or None if at root.
        """
        if len(self._node_stack) >= 2:
            return self._node_stack[-2]

        return None

    def _expr_to_path(self, node: Node) -> Optional[list[str]]:
        """
        Convert an attribute / index chain to a path, resolving the base alias if needed.

        Examples:
        - foo.bar => ["foo", "bar"]
        - foo["bar"].baz => ["foo", "bar", "baz"]
        - foo[i] => ["foo", "<key>"]
        """
        segments: list[str] = []
        current: Optional[Node] = node

        while isinstance(current, (Getattr, Getitem)):
            if isinstance(current, Getattr):
                segments.append(current.attr)
                current = current.node
            else:
                arg = current.arg
                if isinstance(arg, Const) and isinstance(arg.value, str):
                    segments.append(arg.value)
                else:
                    segments.append(DYNAMIC_KEY_PLACEHOLDER)
                current = current.node

        if isinstance(current, Name):
            base = self._resolve_alias([current.name])
            return base + list(reversed(segments))

        return None

    def _iterable_path_from_iter(self, node: Node) -> Optional[list[str]]:
        if self._is_items_call(node):
            assert isinstance(node, Call)
            assert isinstance(node.node, Getattr)
            return self._expr_to_path(node.node.node)

        return self._expr_to_path(node)

    @staticmethod
    def _is_items_call(node: Node) -> bool:
        return isinstance(node, Call) and isinstance(node.node, Getattr) and node.node.attr == "items"

    @staticmethod
    def _target_names(target: Node) -> list[str]:
        """
        Extract loop target variable names, both single (for x in y) and unpacked (for x, y in z).
        """
        if isinstance(target, Name):
            return [target.name]
        if isinstance(target, Tuple):
            return [node.name for node in target.items if isinstance(node, Name)]

        try:
            return [node.name for node in target.find_all(Name)]
        except Exception:
            return []


def collect_variables(jinja_template: str) -> tuple[set[VariablePath], set[IterablePath]]:
    jinja_env = Environment()
    ast = jinja_env.parse(jinja_template)
    collector = VariableCollector()
    collector.visit(ast)

    return collector.paths, collector.iterables
