from __future__ import annotations

from collections.abc import Generator
from typing import Any

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.backend_service import GraphAnalyticsSamplingMode
from graph.exploration.algorithms.graph_features_params import GraphFeaturesParams
from graph.exploration.algorithms.in_memory.graph_features.closeness import apply_closeness
from graph.exploration.algorithms.in_memory.graph_features.clustering import apply_clustering
from graph.exploration.algorithms.in_memory.graph_features.common import (
    build_empty_output_frame,
    build_igraph_graph,
    build_networkx_graph,
    compute_degree_by_key,
    initialize_node_rows,
    iter_edge_batches,
    iter_node_batches,
)
from graph.exploration.algorithms.in_memory.graph_features.connected_components import (
    apply_connected_components,
)
from graph.exploration.algorithms.in_memory.graph_features.degree import apply_degree
from graph.exploration.algorithms.in_memory.graph_features.eigenvector_centrality import (
    apply_eigenvector_centrality,
)
from graph.exploration.algorithms.in_memory.graph_features.pagerank import apply_pagerank
from graph.exploration.algorithms.in_memory.graph_features.square_clustering import (
    apply_square_clustering,
)
from graph.exploration.algorithms.in_memory.graph_features.triangles import apply_triangles
from graph.exploration.algorithms.in_memory.projected_graph_loader import (
    LoadedProjectedGraph,
    load_projected_graph,
)
from graph.exploration.algorithms.projected_graph import ProjectedGraph


def _needs_networkx_graph(params: GraphFeaturesParams) -> bool:
    return any(
        (
            params.degree,
            params.clustering,
            params.triangles and not params.directed_graph,
            params.closeness,
            params.sq_clustering,
            params.connected_components and not params.directed_graph,
        )
    )


def _needs_igraph_graph(params: GraphFeaturesParams) -> bool:
    return params.eigenvector_centrality or params.pagerank


def build_graph_feature_rows(
    graph_data: LoadedProjectedGraph,
    params: GraphFeaturesParams,
) -> list[dict[str, object]]:
    node_rows, rows_by_key = initialize_node_rows(graph_data)
    if _needs_networkx_graph(params):
        graph = build_networkx_graph(graph_data, directed=params.directed_graph)
        if params.degree:
            apply_degree(rows_by_key, compute_degree_by_key(graph_data, directed=params.directed_graph))
        if params.clustering:
            apply_clustering(rows_by_key, graph)
        if not params.directed_graph and params.triangles:
            apply_triangles(rows_by_key, graph)
        if params.closeness:
            apply_closeness(rows_by_key, graph)
        if params.sq_clustering:
            apply_square_clustering(rows_by_key, graph)
        if not params.directed_graph and params.connected_components:
            apply_connected_components(rows_by_key, graph)

    if _needs_igraph_graph(params):
        igraph_graph, ordered_nodes = build_igraph_graph(graph_data, directed=params.directed_graph)
        if params.eigenvector_centrality:
            apply_eigenvector_centrality(
                rows_by_key,
                igraph_graph,
                ordered_nodes,
                directed=params.directed_graph,
            )
        if params.pagerank:
            apply_pagerank(rows_by_key, igraph_graph, ordered_nodes, directed=params.directed_graph)
    return node_rows


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

    if not node_rows:
        yield build_empty_output_frame(params)
        return

    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)
