from __future__ import annotations

import logging
from collections.abc import Generator
from typing import Any

import pandas as pd
from neo4j.exceptions import ClientError

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.projected_graph import 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."""


def _resolve_internal_id(node_dict: dict[str, Any]) -> Any:
    element_id = node_dict.get("_element_id")
    if element_id is not None:
        return element_id
    return node_dict.get("_id")


def ensure_gds_available(conn: DbConnection) -> None:
    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("Neo4j GDS library version: %s", version)
    except Exception as exc:  # pragma: no cover - wrapper behavior tested via callers
        raise Neo4jGDSNotAvailableError("Graph data science plugin is not available.") from exc


def get_node_groups_info(conn: DbConnection, projected_graph: ProjectedGraph) -> dict[str, str]:
    return get_neo4j_node_group_pk_map(conn, projected_graph.node_groups)


def extract_primary_key(node_dict: dict[str, Any], node_groups_info: dict[str, str]) -> Any:
    labels = node_dict.get("_labels", [])
    if not labels:
        return None

    node_label = labels[0]
    primary_key_property = node_groups_info.get(node_label)
    if not primary_key_property:
        return None

    if primary_key_property == "_internal_id":
        return _resolve_internal_id(node_dict)

    return node_dict.get(primary_key_property)


def _quote_cypher_string_literal(value: str) -> str:
    escaped = value.replace("\\", "\\\\").replace("'", "\\'")
    return f"'{escaped}'"


def create_graph_projection(conn: DbConnection, graph_name: str, projected_graph: ProjectedGraph) -> int:
    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(_quote_cypher_string_literal(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(
        {_quote_cypher_string_literal(graph_name)},
        {_quote_cypher_string_literal(node_query)},
        {_quote_cypher_string_literal(relationship_query)}
    )
    YIELD graphName, nodeCount, relationshipCount
    RETURN graphName, nodeCount, relationshipCount
    """

    try:
        result = conn.execute(query)
        df = result.get_as_df()
        if not df.empty:
            return int(df["nodeCount"].iloc[0])
        return 0
    except ClientError as exc:
        if (
            exc.code == "Neo.ClientError.Procedure.ProcedureCallFailed"
            and exc.message
            and "Node-Query returned no nodes" in exc.message
        ):
            return 0
        raise


def drop_graph_projection(conn: DbConnection, graph_name: str) -> None:
    try:
        conn.execute(f"CALL gds.graph.drop('{graph_name}')")
    except ClientError as exc:
        if exc.code == "Neo.ClientError.Procedure.ProcedureCallFailed" and exc.message and "does not exist" in exc.message:
            return
        logger.warning("Failed to drop graph projection '%s'", graph_name, exc_info=True)
    except Exception:
        logger.warning("Failed to drop graph projection '%s'", graph_name, exc_info=True)


def stream_node_metric(
    conn: DbConnection,
    graph_name: str,
    projected_graph: ProjectedGraph,
    *,
    mutate_call: str,
    property_name: str,
    output_column: str,
    batch_size: int = 10000,
) -> Generator[pd.DataFrame, Any, None]:
    ensure_gds_available(conn)
    node_groups_info = get_neo4j_node_group_pk_map(conn, projected_graph.node_groups)

    try:
        total_count = create_graph_projection(conn, graph_name, projected_graph)
        if total_count == 0:
            return

        conn.execute(mutate_call).close()

        offset = 0
        stop = False
        while not stop:
            query = f"""
            CALL gds.graph.nodeProperties.stream('{graph_name}', ['{property_name}'])
            YIELD nodeId, propertyValue
            WITH gds.util.asNode(nodeId) AS node, propertyValue AS {output_column}, nodeId
            RETURN
                {output_column},
                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}
            """
            result = conn.execute(query)
            df = result.get_as_df()
            if df.empty:
                stop = True
                continue

            df["node_id"] = df.apply(
                lambda row: extract_primary_key(
                    {
                        "_labels": row["node_labels"],
                        "_id": row["internal_id"],
                        **(row["node_properties"] if isinstance(row["node_properties"], dict) else {}),
                    },
                    node_groups_info,
                ),
                axis=1,
            )
            df.drop(columns=["internal_id", "node_properties", "node_labels"], inplace=True)
            yield df

            if len(df) < batch_size:
                stop = True
            offset += batch_size
    finally:
        drop_graph_projection(conn, graph_name)
