from __future__ import annotations

from collections import defaultdict
from collections.abc import Hashable

import igraph  # type: ignore[import-untyped]

from graph.exploration.algorithms.graph_clustering_params import ClusterAlgorithm
from graph.exploration.algorithms.in_memory.graph_features.common import (
    FeatureRowsByKey,
    canonical_edge_key,
    initialize_node_rows,
    iter_edge_batches,
    iter_node_batches,
)
from graph.exploration.algorithms.in_memory.projected_graph_loader import LoadedProjectedGraph

ClusterMembership = list[int]


def build_weighted_igraph_graph(
    graph_data: LoadedProjectedGraph,
    *,
    directed: bool,
) -> tuple[igraph.Graph, list[Hashable]]:
    ordered_nodes = list(graph_data.nodes)
    node_index = {node_key: index for index, node_key in enumerate(ordered_nodes)}
    edge_pairs: dict[Hashable, tuple[int, int]] = {}
    aggregated_weights: defaultdict[Hashable, float] = defaultdict(float)

    for edge in graph_data.edges:
        source_index = node_index[edge.source_key]
        target_index = node_index[edge.target_key]
        edge_identity = canonical_edge_key(source_index, target_index, directed=directed)
        edge_pair = (source_index, target_index)
        if not directed and source_index > target_index:
            edge_pair = (target_index, source_index)

        edge_pairs[edge_identity] = edge_pair
        aggregated_weights[edge_identity] += edge.weight if edge.weight is not None else 1.0

    weighted_graph = igraph.Graph(
        n=len(ordered_nodes),
        edges=[edge_pairs[edge_identity] for edge_identity in edge_pairs],
        directed=directed,
    )
    weighted_graph.es["weight"] = [aggregated_weights[edge_identity] for edge_identity in edge_pairs]
    return weighted_graph, ordered_nodes


def fix_dendrogram(graph: igraph.Graph, dendrogram: igraph.clustering.VertexDendrogram) -> None:
    """Patch igraph walktrap dendrograms when disconnected components leave merges incomplete."""
    if not hasattr(dendrogram, "_merges") or not hasattr(dendrogram, "_nmerges"):
        raise RuntimeError(
            "igraph VertexDendrogram internals are not compatible with graph clustering dendrogram repair."
        )
    if not hasattr(dendrogram._merges, "append") or not hasattr(dendrogram._merges, "extend"):
        raise RuntimeError(
            "igraph VertexDendrogram merge storage is not compatible with graph clustering dendrogram repair."
        )

    already_merged: set[int] = set()
    for merge in dendrogram.merges:
        already_merged.update(merge)

    num_dendrogram_nodes = graph.vcount() + len(dendrogram.merges)
    not_merged_yet = sorted(set(range(num_dendrogram_nodes)) - already_merged)
    if len(not_merged_yet) < 2:
        return

    v1, v2 = not_merged_yet[:2]
    # igraph leaves disconnected walktrap dendrograms incomplete, so we patch the
    # protected merge bookkeeping to attach the remaining singleton components.
    try:
        dendrogram._merges.append((v1, v2))
        del not_merged_yet[:2]

        missing_nodes = range(num_dendrogram_nodes, num_dendrogram_nodes + len(not_merged_yet))
        dendrogram._merges.extend(zip(not_merged_yet, missing_nodes))
        dendrogram._nmerges = graph.vcount() - 1
    except (AttributeError, TypeError, ValueError) as error:
        raise RuntimeError(
            "igraph VertexDendrogram internals are not compatible with graph clustering dendrogram repair."
        ) from error


def apply_membership(
    rows_by_key: FeatureRowsByKey,
    ordered_nodes: list[Hashable],
    algorithm: ClusterAlgorithm,
    membership: ClusterMembership,
) -> None:
    for node_key, cluster_id in zip(ordered_nodes, membership):
        rows_by_key[node_key][algorithm] = int(cluster_id)


__all__ = [
    "ClusterAlgorithm",
    "ClusterMembership",
    "FeatureRowsByKey",
    "apply_membership",
    "build_weighted_igraph_graph",
    "fix_dendrogram",
    "initialize_node_rows",
    "iter_edge_batches",
    "iter_node_batches",
]
