from typing import TypedDict, cast

from blends.ctx import ctx
from blends.models import (
    Language,
    NId,
)
from blends.stack.edges import (
    Edge,
    add_edge,
)
from blends.stack.node_helpers import (
    pop_scoped_symbol_node_attributes,
    scope_node_attributes,
)
from blends.stack.scope_types import ScopeType
from blends.syntax.builders.utils import (
    get_next_binding_precedence,
    get_next_synthetic_node_id,
)
from blends.syntax.metadata.java import (
    add_node_range_to_method,
)
from blends.syntax.models import (
    SyntaxGraphArgs,
)


class DirectChildren(TypedDict, total=False):
    block_id: NId
    modifiers_id: NId
    parameters_id: NId


class ListChildren(TypedDict, total=False):
    modifiers_ids: list[NId]
    parameters_ids: list[NId]


def _process_method_children(
    args: SyntaxGraphArgs,
    children: DirectChildren,
) -> None:
    for name_node, child_id in cast("dict[str, NId]", children).items():
        args.syntax_graph.update_node_attribute(args.n_id, name_node, child_id)
        args.syntax_graph.add_edge(
            args.n_id,
            args.generic(args.fork_n_id(child_id)),
            label_ast="AST",
        )


def _process_method_list_children(
    args: SyntaxGraphArgs,
    children: ListChildren,
) -> None:
    for n_ids in cast("dict[str, list[NId]]", children).values():
        for child_id in n_ids:
            args.syntax_graph.add_edge(
                args.n_id,
                args.generic(args.fork_n_id(child_id)),
                label_ast="AST",
            )


def _setup_stack_graph_scope(args: SyntaxGraphArgs, name: str | None) -> None:
    scope_stack = args.metadata.setdefault("scope_stack", [])
    if name and scope_stack and (parent_scope := scope_stack[-1]):
        precedence = get_next_binding_precedence(args)
        method_def_id = get_next_synthetic_node_id(args)
        args.syntax_graph.add_node(
            method_def_id,
            label_type="MethodDeclaration",
            name=name,
            sg_node_type="definition",
        )
        args.syntax_graph.update_node(
            method_def_id,
            pop_scoped_symbol_node_attributes(
                symbol=name,
                precedence=precedence,
            ),
        )
        add_edge(
            args.syntax_graph,
            Edge(source=parent_scope, sink=method_def_id, precedence=precedence),
        )
        add_edge(
            args.syntax_graph,
            Edge(source=method_def_id, sink=args.n_id, precedence=0),
        )
        ast_node = args.ast_graph.nodes.get(args.n_id, {})
        if (label_l := ast_node.get("label_l")) is not None:
            args.syntax_graph.update_node_attribute(method_def_id, "label_l", label_l)
            args.syntax_graph.update_node_attribute(
                method_def_id, "label_c", ast_node.get("label_c", "0")
            )

    args.syntax_graph.update_node(
        args.n_id,
        scope_node_attributes(is_exported=False, scope_type=ScopeType.FUNCTION),
    )
    if scope_stack:
        add_edge(
            args.syntax_graph,
            Edge(source=args.n_id, sink=scope_stack[-1], precedence=0),
        )
    scope_stack.append(args.n_id)


def _cleanup_stack_graph_scope(args: SyntaxGraphArgs) -> None:
    scope_stack = args.metadata.setdefault("scope_stack", [])
    if scope_stack and scope_stack[-1] == args.n_id:
        scope_stack.pop()


def build_method_declaration_node(
    args: SyntaxGraphArgs,
    name: str | None,
    direct_children: DirectChildren,
    list_children: ListChildren,
    access_modifiers: str | None = None,
) -> NId:
    args.syntax_graph.add_node(args.n_id, label_type="MethodDeclaration")
    if name:
        args.syntax_graph.update_node_attribute(args.n_id, "name", name)

    if access_modifiers:
        args.syntax_graph.update_node_attribute(args.n_id, "access_modifiers", access_modifiers)

    if ctx.has_feature_flag("StackGraph"):
        _setup_stack_graph_scope(args, name)

    _process_method_children(args, direct_children)
    _process_method_list_children(args, list_children)

    if args.language == Language.JAVA and args.syntax_graph.nodes.get(NId(0)) and name:
        add_node_range_to_method(args, name)

    if ctx.has_feature_flag("StackGraph"):
        _cleanup_stack_graph_scope(args)

    return args.n_id
