from collections.abc import (
    Callable,
    Iterable,
    Iterator,
)
from itertools import (
    chain,
)

from blends.models import (
    Graph,
    NId,
)

NAttrs = dict[str, str]


def has_labels(n_attrs: NAttrs, **expected_attrs: str) -> bool:
    return all(
        n_attrs.get(expected_attr) == expected_attr_value
        for expected_attr, expected_attr_value in expected_attrs.items()
    )


def predicate_has_labels(**expected_attrs: str) -> Callable[[NAttrs], bool]:
    def predicate(n_attrs: NAttrs) -> bool:
        return has_labels(n_attrs, **expected_attrs)

    return predicate


def filter_nodes(
    graph: Graph,
    nodes: Iterable[NId],
    predicate: Callable[[NAttrs], bool],
) -> tuple[NId, ...]:
    result: tuple[NId, ...] = tuple(n_id for n_id in nodes if predicate(graph.nodes[n_id]))

    return result


def matching_nodes(
    graph: Graph, label_type: str | None = None, **expected_attrs: str
) -> tuple[NId, ...]:
    if label_type:
        candidates = graph.nodes_by_type.get(label_type, [])
        if expected_attrs:
            return filter_nodes(graph, candidates, predicate_has_labels(**expected_attrs))
        return tuple(candidates)
    return filter_nodes(graph, graph.nodes, predicate_has_labels(**expected_attrs))


def _edge_matches(
    graph: Graph,
    source: NId,
    target: NId,
    edge_keys: set[str],
    *,
    strict: bool,
    **edge_attrs: str,
) -> bool:
    if not has_labels(graph[source][target], **edge_attrs):
        return False
    if strict:
        graph_edge_keys = set(graph[source][target].keys())
        difference = graph_edge_keys.difference(edge_keys)
        return not difference or difference == {"label_index"}
    return True


def adj_lazy(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    *,
    strict: bool = False,
    _processed_n_ids: set[NId] | None = None,
    **edge_attrs: str,
) -> Iterator[NId]:
    """Return adjacent nodes to `n_id`, following just edges with given attrs.

    - Parameter `depth` may be -1 to indicate infinite depth.
    - Parameter `strict` indicates that the edges must have only the indicated
      attributes

    - Search is done breadth first.
    - Nodes are returned ordered ascending by index on each level.

    This function must be used instead of graph.adj, because graph.adj
    becomes unstable (unordered) after mutating the graph, also this allow
    following just edges matching `edge_attrs`.
    """
    processed_n_ids: set[NId] = _processed_n_ids or set()
    if depth == 0 or n_id in processed_n_ids:
        return

    processed_n_ids.add(n_id)

    childs: list[NId] = sorted(graph.adj[n_id], key=int)
    edge_keys = set(edge_attrs.keys())
    matching_childs = [
        c_id
        for c_id in childs
        if _edge_matches(graph, n_id, c_id, edge_keys, strict=strict, **edge_attrs)
    ]

    yield from matching_childs

    if depth in {0, 1}:
        return

    for c_id in matching_childs:
        yield from adj_lazy(
            graph,
            c_id,
            depth=depth - 1,
            strict=strict,
            _processed_n_ids=processed_n_ids,
            **edge_attrs,
        )


def adj(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    *,
    strict: bool = False,
    _processed_n_ids: set[NId] | None = None,
    **edge_attrs: str,
) -> tuple[NId, ...]:
    return tuple(
        adj_lazy(
            graph,
            n_id,
            depth=depth,
            strict=strict,
            _processed_n_ids=_processed_n_ids,
            **edge_attrs,
        ),
    )


def adj_ast(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    *,
    strict: bool = False,
    **n_attrs: str,
) -> tuple[NId, ...]:
    return tuple(
        c_id
        for c_id in adj(graph, n_id, depth, strict=strict, label_ast="AST")
        if has_labels(graph.nodes[c_id], **n_attrs)
    )


def adj_cfg(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    *,
    strict: bool = False,
    **n_attrs: str,
) -> tuple[NId, ...]:
    return tuple(
        c_id
        for c_id in adj(graph, n_id, depth, strict=strict, label_cfg="CFG")
        if has_labels(graph.nodes[c_id], **n_attrs)
    )


def adj_cfg_lazy(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    *,
    strict: bool = False,
    **n_attrs: str,
) -> Iterator[NId]:
    yield from adj_lazy(
        graph,
        n_id,
        depth=depth,
        strict=strict,
        _processed_n_ids=set(),
        label_cfg="CFG",
        **n_attrs,
    )


def pred_lazy(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    _processed_n_ids: set[NId] | None = None,
    **edge_attrs: str,
) -> Iterator[NId]:
    """Obtain parent nodes.

    Same as `adj` but follow edges in the opposite direction
    """
    processed_n_ids: set[NId] = _processed_n_ids or set()
    if depth == 0 or n_id in processed_n_ids:
        return

    processed_n_ids.add(n_id)

    p_ids: list[NId] = sorted(graph.pred[n_id], key=int)
    matching_parents = [p_id for p_id in p_ids if has_labels(graph[p_id][n_id], **edge_attrs)]

    yield from matching_parents

    if depth < 0 or depth > 1:
        for p_id in matching_parents:
            yield from pred_lazy(
                graph,
                p_id,
                depth=depth - 1,
                _processed_n_ids=processed_n_ids,
                **edge_attrs,
            )


def pred(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    _processed_n_ids: set[NId] | None = None,
    **edge_attrs: str,
) -> tuple[NId, ...]:
    return tuple(
        pred_lazy(
            graph,
            n_id,
            depth,
            _processed_n_ids=_processed_n_ids,
            **edge_attrs,
        ),
    )


def pred_ast(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    **edge_attrs: str,
) -> tuple[NId, ...]:
    return tuple(pred_ast_lazy(graph, n_id, depth, **edge_attrs))


def pred_ast_lazy(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    **edge_attrs: str,
) -> Iterator[NId]:
    yield from pred_lazy(
        graph,
        n_id,
        depth,
        _processed_n_ids=set(),
        label_ast="AST",
        **edge_attrs,
    )


def pred_cfg(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    **edge_attrs: str,
) -> tuple[NId, ...]:
    return tuple(pred_cfg_lazy(graph, n_id, depth, **edge_attrs))


def pred_cfg_lazy(
    graph: Graph,
    n_id: NId,
    depth: int = 1,
    **edge_attrs: str,
) -> Iterator[NId]:
    yield from pred_lazy(
        graph,
        n_id,
        depth,
        _processed_n_ids=set(),
        label_cfg="CFG",
        **edge_attrs,
    )


def match_ast(
    graph: Graph,
    n_id: NId,
    *label_type: str,
    depth: int = 1,
) -> dict[str, NId | None]:
    index: int = 0
    nodes: dict[str, NId | None] = dict.fromkeys(label_type)

    for c_id in adj_ast(graph, n_id, depth=depth):
        c_type = graph.nodes[c_id]["label_type"]
        if c_type in nodes and nodes[c_type] is None:
            nodes[c_type] = c_id
        else:
            nodes[f"__{index}__"] = c_id
            index += 1

    return nodes


def match_ast_d(graph: Graph, n_id: NId, label_type: str, depth: int = 1) -> NId | None:
    return match_ast(graph, n_id, label_type, depth=depth)[label_type]


def match_ast_group(
    graph: Graph,
    n_id: NId,
    *label_type: str,
    depth: int = 1,
) -> dict[str, list[NId]]:
    index: int = 0
    nodes: dict[str, list[NId]] = {label: [] for label in label_type}

    for c_id in adj_ast(graph, n_id, depth=depth):
        c_type = graph.nodes[c_id]["label_type"]
        if c_type in nodes:
            nodes[c_type].append(c_id)
        else:
            nodes[f"__{index}__"] = [
                c_id,
            ]
            index += 1

    return nodes


def match_ast_group_d(graph: Graph, n_id: NId, label_type: str, depth: int = 1) -> list[NId]:
    return match_ast_group(graph, n_id, label_type, depth=depth)[label_type]


def is_connected_to_cfg(graph: Graph, n_id: NId) -> bool:
    return bool(adj_cfg(graph, n_id) or pred_cfg(graph, n_id))


def lookup_first_cfg_parent(graph: Graph, n_id: NId) -> NId:
    # Lookup first parent who is connected to the CFG
    for p_id in chain([n_id], pred_ast_lazy(graph, n_id, depth=-1)):
        if is_connected_to_cfg(graph, p_id):
            return p_id
    # Base case, pass through
    return n_id


def _get_subgraph(
    graph: Graph,
    node_n_id_predicate: Callable[[NId], bool] = lambda _: True,
    edge_n_attrs_predicate: Callable[[NAttrs], bool] = lambda _: True,
) -> Graph:
    copy: Graph = Graph()

    for n_a_id, n_b_id in graph.edges:
        edge_attrs = graph[n_a_id][n_b_id].copy()
        n_a_attrs = graph.nodes[n_a_id].copy()
        n_b_attrs = graph.nodes[n_b_id].copy()

        if (
            edge_n_attrs_predicate(edge_attrs)
            and node_n_id_predicate(n_a_id)
            and node_n_id_predicate(n_b_id)
        ):
            copy.add_node(n_a_id, **n_a_attrs)
            copy.add_node(n_b_id, **n_b_attrs)
            copy.add_edge(n_a_id, n_b_id, **edge_attrs)

    return copy


def copy_ast(graph: Graph) -> Graph:
    return _get_subgraph(
        graph=graph,
        edge_n_attrs_predicate=predicate_has_labels(label_ast="AST"),
    )


def search_pred_until_type(
    graph: Graph,
    n_id: NId,
    targets: set[str],
    last_child: NId | None = None,
) -> tuple[NId, NId] | None:
    if not last_child:
        last_child = n_id
    if pred_c := pred_ast(graph, n_id):
        if graph.nodes[pred_c[0]]["label_type"] in targets:
            return (pred_c[0], last_child)
        return search_pred_until_type(graph, pred_c[0], targets, pred_c[0])
    return None


def get_nodes_by_path(
    graph: Graph,
    n_id: NId,
    *label_type_path: str,
) -> set[NId]:
    if not label_type_path:
        return set()
    if len(label_type_path) == 1:
        return set(adj_ast(graph, n_id, label_type=label_type_path[0]))

    result: set[NId] = set()
    for node in adj_ast(graph, n_id, label_type=label_type_path[0]):
        result.update(get_nodes_by_path(graph, node, *label_type_path[1:]))
    return result


def get_node_by_path(graph: Graph, n_id: NId, *label_type_path: str) -> NId | None:
    return next(iter(get_nodes_by_path(graph, n_id, *label_type_path)), None)


def nodes_by_type(graph: Graph) -> dict[str, list[NId]]:
    return graph.nodes_by_type
