from __future__ import annotations

import logging

from networkx import MultiDiGraph

from graph.connections.db_connection import DbConnection
from graph.exploration.schema.schema_extractor import NodeTableSchema

logger = logging.getLogger(__name__)


def get_adjacent_nodes(
    conn: DbConnection,
    source_node_schema: NodeTableSchema,
    pk_value: str | int,
    target_node_schema: NodeTableSchema | None,
) -> MultiDiGraph:
    if target_node_schema is None:
        target_string = "dst"
    else:
        target_string = f"dst:`{target_node_schema['label']}`"

    # Get primary key property of source node.
    pk_props = [prop for prop in source_node_schema["properties"] if prop["is_primary_key"]]
    assert len(pk_props) == 1, f"Expected one primary key property, got {pk_props}."
    pk_prop = pk_props[0]

    query_result = None
    try:
        query_result = conn.execute(
            f"""
            MATCH (src:`{source_node_schema["label"]}`)-[r]-({target_string}) WHERE src.`{pk_prop["name"]}` = $pkValue RETURN src, r, dst
            """,
            parameters={"pkValue": pk_value},
        )
        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_node_schema["label"],
            pk_value,
            exc_info=e,
        )
        raise e
    finally:
        if query_result:
            query_result.close()
