from __future__ import annotations

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

import pandas as pd
from neo4j.exceptions import ClientError

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

logger = logging.getLogger(__name__)


class Neo4jGDSNotAvailableError(Exception):
    """Raised when Neo4j GDS library is not available."""

    pass


def ensure_gds_available(conn: DbConnection) -> None:
    """
    Verify GDS library is available in Neo4j.

    Args:
        conn: Neo4j database connection

    Raises:
        Neo4jGDSNotAvailableError: If GDS library is not installed
    """
    try:
        result = conn.execute("RETURN gds.version() AS version")
        df = result.get_as_df()
        if df.empty:
            raise ValueError("GDS version check returned empty result")

        version = df["version"].iloc[0]
        logger.info(f"Neo4j GDS library version: {version}")
    except Exception as e:
        raise Neo4jGDSNotAvailableError("Graph data science plugin is not available.") from e

def extract_primary_key(node_dict: dict, node_groups_info: dict) -> Any:
    """
    Extract primary key value from node properties.

    Args:
        node_dict: Node dictionary from Neo4j result
        node_groups_info: Mapping of node label to primary key property name

    Returns:
        Primary key value
    """
    if not isinstance(node_dict, dict):
        logger.warning(f"Expected dict for node, got {type(node_dict)}: {node_dict}")
        return None

    labels = node_dict.get("_labels", [])
    if not labels:
        logger.warning(f"Node has no labels: {node_dict}")
        return None

    node_label = labels[0]
    primary_key_property = node_groups_info.get(node_label)

    if not primary_key_property:
        logger.warning(f"No primary key property found for label '{node_label}'")
        return None

    if primary_key_property == "_internal_id":
        return node_dict.get("_element_id") or node_dict.get("_id")

    return node_dict.get(primary_key_property)


def create_graph_projection(conn: DbConnection, graph_name: str, projected_graph: ProjectedGraph) -> int:
    """
    Create a GDS graph projection from node/edge groups.

    Args:
        conn: Neo4j database connection
        graph_name: Runtime graph projection name
        projected_graph: Projected graph definition with node and edge groups

    Returns:
        Number of nodes in the projection (0 if empty)

    Raises:
        Exception: If projection creation fails for reasons other than empty graph
    """
    node_label_conditions = " OR ".join([f"n:{label}" for label in projected_graph.node_groups])
    node_query = f"MATCH (n) WHERE {node_label_conditions} RETURN id(n) AS id, labels(n) AS labels"

    edge_types_list = ", ".join([f'"{edge_type}"' for edge_type in projected_graph.edge_groups])
    relationship_query = (
        f"MATCH (s)-[r]->(t) WHERE type(r) IN [{edge_types_list}] "
        f"RETURN id(s) AS source, id(t) AS target, type(r) AS type"
    )

    query = f"""
    CALL gds.graph.project.cypher(
        '{graph_name}',
        '{node_query}',
        '{relationship_query}'
    )
    YIELD graphName, nodeCount, relationshipCount
    RETURN graphName, nodeCount, relationshipCount
    """

    logger.info(f"Creating graph projection '{graph_name}' with query: {query}")

    try:
        result = conn.execute(query)
        df = result.get_as_df()

        if not df.empty:
            node_count = int(df["nodeCount"].iloc[0])
            rel_count = int(df["relationshipCount"].iloc[0])
            logger.info(
                f"Graph projection '{graph_name}' created successfully: "
                f"{node_count} nodes, {rel_count} relationships"
            )
            return node_count
        return 0

    except ClientError as e:
        if (
            e.code == "Neo.ClientError.Procedure.ProcedureCallFailed"
            and e.message
            and "Node-Query returned no nodes" in e.message
        ):
            logger.info(f"Graph projection '{graph_name}' has no nodes - skipping projection")
            return 0
        else:
            logger.exception(f"Failed to create graph projection '{graph_name}'")
            raise
    except Exception as e:
        logger.exception(f"Failed to create graph projection '{graph_name}'")
        raise


def get_node_count_from_projection(conn: DbConnection, graph_name: str) -> int:
    """
    Get total node count from projection for batching.

    Args:
        conn: Neo4j database connection
        graph_name: Name of the graph projection

    Returns:
        Total number of nodes in the projection
    """
    query = f"""
    CALL gds.graph.list('{graph_name}')
    YIELD nodeCount
    RETURN nodeCount
    """
    result = conn.execute(query)
    df = result.get_as_df()

    if df.empty:
        logger.warning(f"Could not get node count for graph '{graph_name}'")
        return 0

    node_count = int(df["nodeCount"].iloc[0])
    logger.debug(f"Graph '{graph_name}' has {node_count} nodes")
    return node_count


def drop_graph_projection(conn: DbConnection, graph_name: str) -> None:
    """
    Drop the graph projection to free memory.

    Args:
        conn: Neo4j database connection
        graph_name: Name of the graph projection to drop
    """
    try:
        query = f"CALL gds.graph.drop('{graph_name}')"
        conn.execute(query)
        logger.info(f"Graph projection '{graph_name}' dropped successfully")
    except ClientError as e:
        if e.code == "Neo.ClientError.Procedure.ProcedureCallFailed" and e.message and "does not exist" in e.message:
            logger.debug(f"Graph projection '{graph_name}' does not exist - already cleaned up")
        else:
            logger.warning(f"Failed to drop graph projection '{graph_name}'", exc_info=True)
    except Exception as e:
        logger.warning(f"Failed to drop graph projection '{graph_name}'", exc_info=True)


def compute_page_rank(
    conn: DbConnection, projected_graph: ProjectedGraph, params: PageRankParams, batch_size: int = 10000
) -> Generator[pd.DataFrame, Any, None]:
    """
    Compute PageRank using Neo4j GDS library with batch processing.

    Args:
        conn: Neo4j database connection
        projected_graph: Projected graph definition with node and edge groups
        params: PageRank algorithm parameters
        batch_size: Number of results to process per batch

    Yields:
        DataFrames containing columns: rank, node_group, node_id

    Raises:
        Neo4jGDSNotAvailableError: If GDS library is not installed
    """
    ensure_gds_available(conn)

    runtime_graph_name = f"page_rank_graph_{uuid4().hex}"
    node_groups_info = get_neo4j_node_group_pk_map(conn, projected_graph.node_groups)

    try:
        total_count = create_graph_projection(conn, runtime_graph_name, projected_graph)

        if total_count == 0:
            logger.info(f"Graph projection '{runtime_graph_name}' contains no nodes - no PageRank to compute")
            return

        if params.normalizeInitial:
            logger.warning(
                "Neo4j GDS does not support normalizeInitial parameter. Results will not be normalized to sum to 1."
            )

        def extract_pk(row):
            node_dict = {
                "_labels": row["node_labels"],
                "_id": row["internal_id"],
                **(row["node_properties"] if isinstance(row["node_properties"], dict) else {}),
            }
            return extract_primary_key(node_dict, node_groups_info)

        offset = 0
        stop = False

        while not stop:
            query = f"""
            CALL gds.pageRank.stream('{runtime_graph_name}', {{
                dampingFactor: {params.damping_factor},
                maxIterations: {params.max_iterations},
                tolerance: {params.tolerance}
            }})
            YIELD nodeId, score
            WITH gds.util.asNode(nodeId) AS node, score AS rank, nodeId
            RETURN
                rank,
                labels(node)[0] AS node_group,
                nodeId AS internal_id,
                properties(node) AS node_properties,
                labels(node) AS node_labels
            ORDER BY internal_id
            SKIP {offset}
            LIMIT {batch_size}
            """

            logger.debug(f"Executing PageRank query with offset {offset}, limit {batch_size}")
            result = conn.execute(query)
            df = result.get_as_df()

            if df.empty:
                logger.info("Batch is empty, no more results to iterate over.")
                stop = True
            else:
                logger.debug(f"Processing results of length {len(df)}.")

                df["node_id"] = df.apply(extract_pk, axis=1)

                df.drop(columns=["internal_id", "node_properties", "node_labels"], 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:
        drop_graph_projection(conn, runtime_graph_name)
