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

from blends.models import (
    Graph,
    NId,
    SyntaxGraph,
)

NAttrs = Mapping[str, str | NId]


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


def get_args_n_ids(graph: Graph, n_id: NId) -> list[NId]:
    if (args_parent_n_id := graph.nodes[n_id].get("arguments_id")) is None:
        return []

    return [
        arg_id
        for arg_id in adj_ast(graph, args_parent_n_id)
        if graph.nodes[arg_id].get("label_type") != "Comment"
    ]


def get_n_arg(graph: Graph, n_id: NId | None, arg_idx: int) -> NId | None:
    if not n_id:
        return None

    for idx, arg_id in enumerate(get_args_n_ids(graph, n_id)):
        if idx == arg_idx:
            return arg_id
    return None


def resolve_symbol_scope_index(graph: SyntaxGraph, n_id: NId) -> NId | None:
    symbol = graph.nodes[n_id].get("symbol")
    scope: NId | None = graph.nodes[n_id].get("symbol_scope")
    if not symbol or scope is None:
        return None
    defs = graph.symbol_index.get(symbol, {})
    while scope is not None:
        if def_nid := defs.get(scope):
            return def_nid
        scope = graph.scope_parent.get(scope)
    return None


def _literal_binary_operation(graph: SyntaxGraph, n_id: NId, level: int) -> str | None:
    nodes = graph.nodes
    if left_n_id := nodes[n_id].get("left_id"):
        if nodes[left_n_id].get("label_type") == "Literal":
            return nodes[left_n_id].get("value")
        if nodes[left_n_id].get("label_type") == "BinaryOperation":
            return _literal_binary_operation(graph, left_n_id, level + 1)
    if (
        (right_n_id := nodes[n_id].get("right_id"))
        and level == 0
        and nodes[right_n_id].get("label_type") == "Literal"
    ):
        return nodes[right_n_id].get("value")
    return None


def _resolve_child_symbols(graph: SyntaxGraph, n_id: NId, depth: int) -> str | None:
    nodes = graph.nodes
    for child_id in match_ast_group_d(graph, n_id, "SymbolLookup", -1):
        if (parent := next(iter(pred_ast(graph, child_id)), None)) and (
            nodes[parent].get("label_type") == "ArgumentList"
        ):
            continue
        def_nid = resolve_symbol_scope_index(graph, child_id)
        if def_nid is not None and (result := _resolve_value(graph, def_nid, depth)):
            return result
    return None


def _check_direct_string(graph: SyntaxGraph, n_id: NId) -> str | None:
    nodes = graph.nodes
    label_type = nodes[n_id].get("label_type")

    if (
        label_type == "Literal"
        and nodes[n_id].get("value_type") != "null"
        and len(adj_ast(graph, n_id)) == 0
    ):
        return nodes[n_id].get("value")

    if (obj_id := nodes[n_id].get("object_id")) and (nodes[obj_id].get("label_type") == "Literal"):
        return nodes[obj_id].get("value")

    if label_type == "BinaryOperation":
        return _literal_binary_operation(graph, n_id, 0)

    return None


def _resolve_value(graph: SyntaxGraph, n_id: NId, depth: int) -> str | None:
    value_id = graph.nodes[n_id].get("value_id")
    if not value_id:
        return None
    return has_string_definition_resolver(graph, value_id, _depth=depth)


# WIP: Only works when the ScopeIndex feature flag is active.
# Do not use in production methods until ScopeIndex is fully rolled out.
def has_string_definition_resolver(
    graph: SyntaxGraph, n_id: NId | None, *, _depth: int = 0
) -> str | None:
    """Resolve whether a node traces back to a string literal.

    Uses ScopeIndex for symbol resolution instead of CFG backward paths.
    Recurses through SymbolLookup chains up to max_resolver_depth.
    """
    max_resolver_depth = 4
    if not n_id or _depth > max_resolver_depth:
        return None

    if result := _check_direct_string(graph, n_id):
        return result

    if graph.nodes[n_id].get("label_type") == "SymbolLookup":
        def_nid = resolve_symbol_scope_index(graph, n_id)
        if def_nid is None:
            return None
        return _resolve_value(graph, def_nid, _depth + 1)

    return _resolve_child_symbols(graph, n_id, _depth + 1)
