from blends.models import (
    NId,
)
from blends.query import (
    adj_ast,
)
from blends.stack.attributes import (
    STACK_GRAPH_IS_REFERENCE,
)
from blends.stack.edges import (
    Edge,
    add_edge,
)
from blends.stack.node_helpers import (
    pop_scoped_symbol_node_attributes,
    root_node_attributes,
)
from blends.stack.policies.registry import (
    get_export_policy,
)
from blends.syntax.models import (
    SyntaxGraphArgs,
)
from blends.utilities.text_nodes import (
    node_to_str,
)

_PATTERN_NODE_TYPES: frozenset[str] = frozenset(
    {"pattern_list", "tuple_pattern", "list_pattern", "list_splat_pattern"}
)


def extract_pattern_identifier_symbols(
    args: SyntaxGraphArgs, *, pattern_id: NId
) -> list[tuple[NId, str]]:
    label_type = args.ast_graph.nodes.get(pattern_id, {}).get("label_type")
    if label_type == "identifier":
        symbol = node_to_str(args.ast_graph, pattern_id)
        return [(pattern_id, symbol)] if symbol else []
    if label_type in _PATTERN_NODE_TYPES:
        return [
            result
            for child_id in args.ast_graph.successors(pattern_id)
            for result in extract_pattern_identifier_symbols(args, pattern_id=child_id)
        ]
    return []


def add_pattern_binding_nodes(args: SyntaxGraphArgs, *, var_id: NId) -> None:
    scope_stack = args.metadata.setdefault("scope_stack", [])
    parent_scope = scope_stack[-1] if scope_stack else None
    for identifier_nid, symbol in extract_pattern_identifier_symbols(args, pattern_id=var_id):
        synthetic_nid = get_next_synthetic_node_id(args)
        precedence = get_next_binding_precedence(args)
        stack_attrs = pop_scoped_symbol_node_attributes(symbol=symbol, precedence=precedence)
        if args.stack_graph is not None:
            args.stack_graph.add_node(synthetic_nid, **stack_attrs)
        if parent_scope:
            edge = Edge(source=parent_scope, sink=synthetic_nid, precedence=precedence)
            if args.stack_graph is not None:
                add_edge(args.stack_graph, edge)
        if args.syntax_graph.nodes.get(identifier_nid) and args.stack_graph is not None:
            ref_attr: dict[str, object] = {STACK_GRAPH_IS_REFERENCE: False}
            args.stack_graph.add_node(identifier_nid, **ref_attr)


def register_multi_binding(args: SyntaxGraphArgs, var_id: NId, scope: NId) -> None:
    for child in adj_ast(args.syntax_graph, var_id, label_type="SymbolLookup"):
        if (sym := args.syntax_graph.nodes[child].get("symbol")) and sym != "_":
            args.syntax_graph.symbol_index.setdefault(sym, {}).setdefault(scope, []).append(
                args.n_id
            )


def setup_scope_index_binding(args: SyntaxGraphArgs, var_id: NId) -> None:
    if not (scope_stack := args.metadata.get("scope_stack")):
        return
    if (symbol := bound_identifier_symbol(args, var_id=var_id)) is not None:
        args.syntax_graph.symbol_index.setdefault(symbol, {}).setdefault(
            scope_stack[-1], []
        ).append(args.n_id)
    elif args.syntax_graph.nodes.get(var_id, {}).get("label_type") == "ArgumentList":
        register_multi_binding(args, var_id, scope_stack[-1])


def bound_identifier_symbol(args: SyntaxGraphArgs, *, var_id: NId) -> str | None:
    var_node = args.ast_graph.nodes.get(var_id, {})
    if var_node.get("label_type") != "identifier":
        return None
    symbol = node_to_str(args.ast_graph, var_id)
    return symbol or None


def get_next_binding_precedence(args: SyntaxGraphArgs) -> int:
    scope_stack = args.metadata.get("scope_stack", [])
    if not scope_stack:
        return 0
    return 1


def is_definition_exported(args: SyntaxGraphArgs) -> bool:
    policy = get_export_policy(args.language)
    scope_stack = args.metadata.get("scope_stack", [])
    parent_scope = scope_stack[-1] if scope_stack else None
    return policy.is_definition_exported(
        args.syntax_graph, args.stack_graph, args.n_id, parent_scope
    )


def get_next_synthetic_node_id(args: SyntaxGraphArgs) -> NId:
    current_id = args.metadata.get("next_synthetic_id")
    if current_id is None:
        msg = "next_synthetic_id not initialized in metadata"
        raise ValueError(msg)
    args.metadata["next_synthetic_id"] = current_id + 1
    return current_id


def get_or_create_root_nid(args: SyntaxGraphArgs) -> NId:
    existing = args.metadata.get("root_nid")
    if existing is not None:
        return existing
    root_nid = get_next_synthetic_node_id(args)
    attrs = root_node_attributes()
    args.syntax_graph.add_node(root_nid)
    if args.stack_graph is not None:
        args.stack_graph.add_node(root_nid, **attrs)
    args.metadata["root_nid"] = root_nid
    return root_nid
