from __future__ import annotations

import logging
from typing import List

from networkx import MultiDiGraph
from pandas import DataFrame

from graph.connections.db_connection import DbConnection
from graph.exploration.queries.node_adjacency import AdjacentNodeGroupInfo
from graph.exploration.schema.schema_extractor import NodeTableSchema

logger = logging.getLogger(__name__)


def get_adjacent_node_groups_info(conn: DbConnection, node_schema: NodeTableSchema, node_id: str | int) -> List[AdjacentNodeGroupInfo]:
    """Get information about adjacent nodes for a given node.
    
    This function queries the Neo4j database to find all nodes connected to a specified node,
    groups them by label, and returns the count of nodes in each group. It handles nodes that
    were created both inside and outside of Visual Graph by trying to match using both
    elementId and primary key properties.
    
    Args:
        conn: Active database connection to execute the query.
        node_schema: Schema information for the source node including label and properties.
        node_id: The identifier of the source node. Can be either an elementId (string)
                or a primary key value (string or int).
    
    Returns:
        A list of dictionaries containing the label and count of adjacent nodes for each
        target node type. Example: [{"label": "Person", "count": 5}, {"label": "Company", "count": 2}]
    
    Note:
        - Falls back to "id" as the primary key property name if no primary key is found in the schema.
        - The WHERE clause checks both primary key match and elementId match to handle graphs
          created inside or outside of Visual Graph.
    """
    pk_props = [prop for prop in node_schema["properties"] if prop["is_primary_key"]]
    # pk_prop may not exists if the graph has been created outside of VG
    try:
        pk_prop = pk_props[0]["name"]
    except IndexError:
        logger.warning("No primary key found for node in the schema.")
        pk_prop = "id"

    label = node_schema["label"]

    # Query to handle a nodeId that has been created inside or outside of VG
    cypher_query = f"""
        MATCH (n:`{label}`)
        WHERE n.`{pk_prop}` = $nodeId OR elementId(n) = $nodeId
        MATCH (n)-[]-(t)
        UNWIND labels(t) AS targetLabel
        RETURN targetLabel, COUNT(DISTINCT t) AS count
    """

    query_result = conn.execute(
        cypher_query,
        parameters={"nodeId": node_id}
    )

    df: DataFrame = query_result.get_as_df()

    return [{"label": row["targetLabel"], "count": row["count"]} for _, row in df.iterrows()]


def get_adjacent_nodes(
    conn: DbConnection,
    source_node_schema: NodeTableSchema,
    node_id: str | int,
    target_node_schema: NodeTableSchema | None,
) -> MultiDiGraph:
    """Get all adjacent nodes and their connecting relationships as a NetworkX graph.
    
    This function retrieves all nodes connected to a specified source node, optionally filtered
    by target node type. It returns the complete subgraph including source node, target nodes,
    and all relationships between them. The function handles nodes created both inside and
    outside Visual Graph by matching on either primary key or elementId.
    
    Args:
        conn: Active database connection to execute the query.
        source_node_schema: Schema information for the source node including label and properties.
        node_id: The identifier of the source node. Can be either an elementId (string)
                or a primary key value (string or int).
        target_node_schema: Optional schema information for filtering target nodes by label.
                           If None, all connected nodes regardless of type are returned.
    
    Returns:
        A NetworkX MultiDiGraph containing the source node, all adjacent nodes matching the
        target criteria, and the relationships connecting them. Nodes and edges include
        all their properties from the database.
    
    Raises:
        RuntimeError: If there is an error executing the query or retrieving the graph data.
                     The error is logged before being re-raised.
    
    Note:
        - Falls back to "id" as the primary key property name if no primary key is found in the schema.
        - The WHERE clause checks both primary key match and elementId match to handle graphs
          created inside or outside of Visual Graph.
    """
    if target_node_schema is None:
        target_string = "dst"
    else:
        target_string = f"dst:`{target_node_schema['label']}`"

    source_label = source_node_schema["label"]

    # pk_prop may not exists if the graph has been created outside of VG
    pk_props = [prop for prop in source_node_schema["properties"] if prop["is_primary_key"]]
    try:
        pk_prop = pk_props[0]["name"]
    except IndexError:
        logger.warning("No primary key found for node in the schema.")
        pk_prop = "id"

    try:
        query_result = conn.execute(
            f"""
            MATCH (src:`{source_label}`)-[r]-({target_string})
            WHERE src.`{pk_prop}` = $nodeId OR elementId(src) = $nodeId
            RETURN src, r, dst
            """,
            parameters={"nodeId": node_id},
        )
        return query_result.get_as_networkx()
    except RuntimeError as e:
        logger.error(
            f"Error while getting adjacent nodes for node %s with PK=%s.",
            source_label,
            node_id,
            exc_info=e,
        )
        raise e
