from blends.ctx import ctx
from blends.models import (
    Graph,
    NId,
)
from blends.query import (
    adj_ast,
    match_ast_d,
)
from blends.syntax.builders.parameter import (
    build_parameter_node,
)
from blends.syntax.models import (
    SyntaxGraphArgs,
)
from blends.utilities.text_nodes import (
    node_to_str,
)

_NONE_NAMES: frozenset[str] = frozenset({"None", "none"})


def _name_from_leaf_node(
    graph: Graph,
    type_id: NId,
    node_type: str,
) -> set[str]:
    node = graph.nodes[type_id]
    if node_type == "identifier":
        name = node.get("label_text", "")
        return {name} if name and name not in _NONE_NAMES else set()
    if node_type == "string":
        name = node_to_str(graph, type_id).strip("\"'").split(".")[-1]
        return {name} if name and name not in _NONE_NAMES else set()
    attr_id = node.get("label_field_attribute")
    return _extract_annotation_class_names(graph, attr_id) if attr_id else set()


def _collect_generic_type_names(graph: Graph, type_id: NId) -> set[str]:
    type_param_id = match_ast_d(graph, type_id, "type_parameter")
    if not type_param_id:
        return set()
    names: set[str] = set()
    for child_id in adj_ast(graph, type_param_id):
        if graph.nodes.get(child_id, {}).get("label_type", "") in {",", "[", "]"}:
            continue
        names |= _extract_annotation_class_names(graph, child_id)
    return names - _NONE_NAMES


def _collect_union_names(graph: Graph, type_id: NId) -> set[str]:
    names: set[str] = set()
    for child_id in adj_ast(graph, type_id):
        if graph.nodes.get(child_id, {}).get("label_type", "") == "|":
            continue
        names |= _extract_annotation_class_names(graph, child_id)
    return names - _NONE_NAMES


_LEAF_NODE_TYPES: frozenset[str] = frozenset({"identifier", "string", "attribute"})
_UNION_NODE_TYPES: frozenset[str] = frozenset({"binary_operator", "union_type"})


def _extract_annotation_class_names(graph: Graph, type_id: NId) -> set[str]:
    node: dict[str, str] = graph.nodes.get(type_id, {})
    node_type = node.get("label_type", "")

    if node_type == "type":
        children = adj_ast(graph, type_id)
        return _extract_annotation_class_names(graph, children[0]) if children else set()

    if node_type in _LEAF_NODE_TYPES:
        return _name_from_leaf_node(graph, type_id, node_type)

    if node_type == "generic_type":
        return _collect_generic_type_names(graph, type_id)

    if node_type in _UNION_NODE_TYPES:
        return _collect_union_names(graph, type_id)

    return set()


def reader(args: SyntaxGraphArgs) -> NId:
    graph = args.ast_graph
    n_attrs = graph.nodes[args.n_id]
    var_type = None
    var_name = "UnnamedVar"
    annotation_class_names: set[str] | None = None

    type_id = n_attrs.get("label_field_type")
    if type_id:
        var_type = node_to_str(graph, type_id)
        if ctx.has_feature_flag("StackGraph"):
            annotation_class_names = _extract_annotation_class_names(graph, type_id)

    identifier_id = n_attrs.get("label_field_name") or match_ast_d(graph, args.n_id, "identifier")
    if identifier_id:
        var_name = node_to_str(graph, identifier_id)

    value_id = n_attrs.get("label_field_value")
    return build_parameter_node(
        args=args,
        variable=var_name,
        variable_type=var_type,
        value_id=value_id,
        annotation_class_names=annotation_class_names,
    )
