from __future__ import annotations

import logging
from typing import List

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, pk_value: str | int
) -> List[AdjacentNodeGroupInfo]:
    pk_props = [prop for prop in 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 (n:`{node_schema["label"]}` {{`{pk_prop["name"]}`: $pk_value}})-[]-(t) 
            RETURN LABEL(t) AS targetLabel, COUNT(DISTINCT t) AS count
            """,
            parameters={"pk_value": pk_value},
        )
        df: DataFrame = query_result.get_as_df()

        return [{"label": row["targetLabel"], "count": row["count"]} for _, row in df.iterrows()]
    except RuntimeError as e:
        logger.error(
            f"Error while getting adjacent node info for node %s with PK=%s.",
            node_schema["label"],
            pk_value,
            exc_info=e,
        )
        return []
    finally:
        if query_result:
            query_result.close()
