from __future__ import annotations

import logging
from collections.abc import Hashable
from typing import Any, 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.in_memory.projected_graph_loader import (
    LoadedProjectedGraph,
    load_projected_graph,
)
from graph.exploration.algorithms.page_rank import PageRankParams, ProjectedGraph

logger = logging.getLogger(__name__)


def _compute_page_rank_scores(
    graph_data: LoadedProjectedGraph,
    params: PageRankParams,
) -> dict[Hashable, float]:
    if not graph_data.nodes:
        return {}

    ordered_nodes = list(graph_data.nodes)
    index_by_node = {node_key: index for index, node_key in enumerate(ordered_nodes)}
    graph = igraph.Graph(
        n=len(ordered_nodes),
        edges=[(index_by_node[edge.source_key], index_by_node[edge.target_key]) for edge in graph_data.edges],
        directed=True,
    )
    graph.vs["name"] = [str(node_key) for node_key in ordered_nodes]
    try:
        raw_scores = graph.pagerank(
            directed=True,
            damping=params.damping_factor,
            weights=None,
            implementation="prpack",
        )
    except igraph.InternalError:
        logger.warning("PRPACK PageRank failed; retrying with igraph default implementation.", exc_info=True)
        raw_scores = graph.pagerank(
            directed=True,
            damping=params.damping_factor,
            weights=None,
        )

    return {node_key: float(raw_scores[index]) for index, node_key in enumerate(ordered_nodes)}


def _iter_page_rank_rows(
    graph_data: LoadedProjectedGraph,
    scores: dict[Hashable, float],
    batch_size: int,
) -> Generator[pd.DataFrame, Any, None]:
    rows = []
    for node_key, node_data in graph_data.nodes.items():
        rows.append(
            {
                "rank": scores.get(node_key, 0.0),
                "node_group": node_data.node_group,
                "node_id": node_data.node_id,
            }
        )

    rows.sort(key=lambda row: (str(row["node_group"]), str(row["node_id"])))

    for offset in range(0, len(rows), batch_size):
        yield pd.DataFrame(rows[offset : offset + batch_size])


def compute_page_rank(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    params: PageRankParams,
    batch_size: int = 10000,
    *,
    sampling_mode: GraphAnalyticsSamplingMode = GraphAnalyticsSamplingMode.NONE,
    sampling_limit: int | None = None,
) -> Generator[pd.DataFrame, Any, None]:
    if params.normalizeInitial:
        logger.warning("In-memory PageRank ignores normalizeInitial because igraph does not expose an equivalent option.")
    if GraphAnalyticsSamplingMode(sampling_mode) != GraphAnalyticsSamplingMode.NONE:
        logger.warning(
            "In-memory PageRank sampling does not guarantee today that we keep edges for the nodes we selected. "
            "Results can therefore be sparse, edgeless, or a partial view of the real PageRank neighborhood."
        )
    graph_data = load_projected_graph(
        conn,
        projected_graph,
        sampling_mode=sampling_mode,
        sampling_limit=sampling_limit,
    )
    scores = _compute_page_rank_scores(graph_data, params)
    yield from _iter_page_rank_rows(graph_data, scores, batch_size)
