import logging
from typing import Any, Generator
from uuid import uuid4

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.page_rank import PageRankParams, ProjectedGraph
from graph.exploration.schema.kuzu.schema_extractor import get_node_group_pk_map as get_kuzu_node_group_pk_map

logger = logging.getLogger(__name__)


def _extract_node_id(node: object, node_groups_info: dict[str, str]) -> Any:
    if not isinstance(node, dict):
        return pd.NA

    node_label = node.get("_label")
    if not isinstance(node_label, str):
        return pd.NA

    primary_key = node_groups_info.get(node_label)
    if primary_key is None:
        return pd.NA

    return node.get(primary_key, pd.NA)


def create_projected_graph(conn: DbConnection, graph_name: str, projected_graph: ProjectedGraph) -> None:
    """
    Create a projected graph in Kuzu for algorithm execution.

    Args:
        conn: Kuzu database connection
        graph_name: Runtime graph projection name
        projected_graph: Projected graph definition with node groups and edge groups
    """
    node_groups_str = ", ".join(f"'{group}'" for group in projected_graph.node_groups)
    edge_groups_str = ", ".join(f"'{group}'" for group in projected_graph.edge_groups)

    query = f"""
    CALL project_graph(
        '{graph_name}',
        [{node_groups_str}],
        [{edge_groups_str}]
    )
    """
    logger.info(f"Creating projected graph with query: {query}.")
    conn.execute(query)

def compute_page_rank(
    conn: DbConnection, projected_graph: ProjectedGraph, params: PageRankParams, batch_size: int = 10000
) -> Generator[pd.DataFrame, Any, None]:
    """
    Returns:
        a generator that yields DataFrames containing the columns: "rank", "node_group", and "node_id".
    """
    runtime_graph_name = f"page_rank_graph_{uuid4().hex}"

    create_projected_graph(conn, runtime_graph_name, projected_graph)

    node_groups_info = get_kuzu_node_group_pk_map(conn)

    offset = 0
    stop = False
    while not stop:
        query = f"""
        CALL page_rank(
            '{runtime_graph_name}',
            dampingFactor := {params.damping_factor},
            maxIterations := {params.max_iterations},
            tolerance := {params.tolerance},
            normalizeInitial := {params.normalizeInitial}
        )
        RETURN node, rank, OFFSET(ID(node)) AS internal_offset
        ORDER BY internal_offset DESC
        SKIP {offset}
        LIMIT {batch_size};
        """
        logger.debug(f"Executing query {query}.")
        result = conn.execute(query)
        try:
            df = result.get_as_df()
            if df.empty:
                logger.info("Batch is empty, not more results to iterate over.")
                stop = True
            else:
                logger.debug(f"Processing results of length {len(df)}.")

                df["node_group"] = df["node"].apply(lambda x: x.get("_label") if isinstance(x, dict) else pd.NA)
                df["node_id"] = df["node"].apply(lambda x: _extract_node_id(x, node_groups_info))

                df.drop(columns=["node"], inplace=True)
                df.drop(columns=["internal_offset"], inplace=True)

                yield df

                if len(df) < batch_size:
                    logger.info("Batch is below max batch size, no more results to iterate over.")
                    stop = True
                offset += batch_size
        finally:
            result.close()
