import logging
from typing import Any, Generator

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.page_rank import PageRankParams, ProjectedGraph

logger = logging.getLogger(__name__)


def install_algo_extension(conn: DbConnection) -> None:
    conn.execute("INSTALL ALGO; LOAD ALGO;")
    logger.info("Kuzu algorithm extension installed successfully.")


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

    Args:
        conn: Kuzu database connection
        projected_graph: Projected graph definition with name, 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(
        '{projected_graph.name}',
        [{node_groups_str}],
        [{edge_groups_str}]
    )
    """
    logger.info(f"Creating projected graph with query: {query}.")
    conn.execute(query)


def __get_node_groups_info__(conn: DbConnection) -> dict[str, str]:
    groups_name_df: pd.DataFrame = conn.execute("CALL show_tables() WHERE type = 'NODE' RETURN name;").get_as_df()  # type: ignore
    result = {}
    for _, row in groups_name_df.iterrows():
        group_name = row["name"]  # type: ignore
        groups_info_df = conn.execute(f'CALL TABLE_INFO("{group_name}") RETURN *;').get_as_df()  # type: ignore
        pk_name = groups_info_df[groups_info_df["primary key"] == True].iloc[0]["name"]  # noqa: E712

        result[group_name] = pk_name

    return result


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".
    """
    install_algo_extension(conn)
    create_projected_graph(conn, projected_graph)

    node_groups_info = __get_node_groups_info__(conn)

    offset = 0
    stop = False
    while not stop:
        query = f"""
        CALL page_rank(
            '{projected_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: x.get(node_groups_info.get(x.get("_label"))) if isinstance(x, dict) else pd.NA  # type: ignore
                )

                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()
