from __future__ import annotations

from collections.abc import Generator

import igraph  # type: ignore[import-untyped]
import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.backend_service import GraphAnalyticsSamplingMode
from graph.exploration.algorithms.graph_clustering_params import (
    ClusterAlgorithm,
    GraphClusteringParams,
    selected_graph_clustering_names,
)
from graph.exploration.algorithms.in_memory.graph_clustering.common import (
    apply_membership,
    build_weighted_igraph_graph,
    fix_dendrogram,
    initialize_node_rows,
    iter_edge_batches,
    iter_node_batches,
)
from graph.exploration.algorithms.in_memory.projected_graph_loader import (
    LoadedProjectedGraph,
    load_projected_graph,
)
from graph.exploration.algorithms.projected_graph import ProjectedGraph

INFOMAP_TRIALS = 10

CLUSTERING_ALGORITHMS: dict[ClusterAlgorithm, tuple[str, dict[str, object]]] = {
    "fastgreedy": ("community_fastgreedy", {"weights": "weight"}),
    "multilevel": ("community_multilevel", {"weights": "weight"}),
    # Keep a fixed Infomap trial count in phase 1 to bound runtime without expanding the UI surface.
    "infomap": ("community_infomap", {"edge_weights": "weight", "trials": INFOMAP_TRIALS}),
    "walktrap": ("community_walktrap", {"weights": "weight", "steps": 10}),
}
def build_graph_clustering_rows(
    graph_data: LoadedProjectedGraph,
    params: GraphClusteringParams,
) -> list[dict[str, object]]:
    graph, ordered_nodes = build_weighted_igraph_graph(
        graph_data,
        directed=params.directed_graph,
    )
    node_rows, rows_by_key = initialize_node_rows(graph_data)

    for algorithm in selected_graph_clustering_names(params):
        method_name, method_kwargs = CLUSTERING_ALGORITHMS[algorithm]
        result = getattr(graph, method_name)(**method_kwargs)
        membership = to_membership(graph, result)
        apply_membership(rows_by_key, ordered_nodes, algorithm, membership)

    return node_rows


def to_membership(
    graph: igraph.Graph,
    result: igraph.clustering.VertexDendrogram | igraph.clustering.VertexClustering,
) -> list[int]:
    if isinstance(result, igraph.clustering.VertexDendrogram):
        fix_dendrogram(graph, result)
        return [int(cluster_id) for cluster_id in result.as_clustering().membership]
    if isinstance(result, igraph.clustering.VertexClustering):
        return [int(cluster_id) for cluster_id in result.membership]
    raise ValueError("Unsupported igraph clustering result type.")


def compute_graph_clustering(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    params: GraphClusteringParams,
    *,
    batch_size: int = 10000,
    sampling_mode: GraphAnalyticsSamplingMode = GraphAnalyticsSamplingMode.NONE,
    sampling_limit: int | None = None,
) -> Generator[pd.DataFrame, None, None]:
    graph_data = load_projected_graph(
        conn,
        projected_graph,
        sampling_mode=sampling_mode,
        sampling_limit=sampling_limit,
        weight_property=params.weight_property,
    )
    node_rows = build_graph_clustering_rows(graph_data, params)

    if params.output_type == "output_edges":
        yield from iter_edge_batches(graph_data, node_rows, batch_size)
        return

    yield from iter_node_batches(node_rows, batch_size)
