from __future__ import annotations

from typing import List, TypedDict

from networkx import MultiDiGraph

from graph.connections.db_connection import DbConnection
from graph.connections.kuzu.db_connection import KuzuConnection
from graph.connections.neo4j.db_connection import Neo4jConnection
from graph.exploration.schema.schema_extractor import NodeTableSchema


class AdjacentNodeGroupInfo(TypedDict):
    label: str
    count: int


def get_adjacent_node_groups_info(
    conn: DbConnection, node_schema: NodeTableSchema, pk_value: str | int
) -> List[AdjacentNodeGroupInfo]:
    from graph.exploration.queries.kuzu.get_adjacent_node_groups_info import (
        get_adjacent_node_groups_info as get_adjacent_node_groups_info_kuzu,
    )
    from graph.exploration.queries.neo4j.node_adjacency import (
        get_adjacent_node_groups_info as get_adjacent_node_groups_info_neo4j,
    )

    if isinstance(conn, KuzuConnection):
        return get_adjacent_node_groups_info_kuzu(conn, node_schema, pk_value)

    elif isinstance(conn, Neo4jConnection):
        return get_adjacent_node_groups_info_neo4j(conn, node_schema, pk_value)
    else:
        raise ValueError(f"Unsupported connection type: {type(conn)}")


def get_adjacent_nodes(
    conn: DbConnection,
    source_node_schema: NodeTableSchema,
    pk_value: str | int,
    target_node_schema: NodeTableSchema | None,
) -> MultiDiGraph:
    from graph.exploration.queries.kuzu.get_adjacent_nodes import (
        get_adjacent_nodes as get_adjacent_nodes_kuzu,
    )
    from graph.exploration.queries.neo4j.node_adjacency import (
        get_adjacent_nodes as get_adjacent_nodes_neo4j,
    )

    if isinstance(conn, KuzuConnection):
        return get_adjacent_nodes_kuzu(conn, source_node_schema, pk_value, target_node_schema)

    elif isinstance(conn, Neo4jConnection):
        return get_adjacent_nodes_neo4j(conn, source_node_schema, pk_value, target_node_schema)

    else:
        raise ValueError(f"Unsupported connection type: {type(conn)}")
