from typing import Dict, TypedDict

from typing_extensions import TypeAlias

from graph.connections.db_connection import DbConnection

GroupCounts: TypeAlias = Dict[str, int]


class GraphTotalCounts(TypedDict):
    nodes: GroupCounts
    edges: GroupCounts


def get_node_counts_by_group(conn: DbConnection) -> GroupCounts:
    """
    Get count of nodes grouped by label.
    For nodes with multiple labels, each label is counted separately.

    Args:
        conn: Database connection

    Returns:
        Dictionary mapping node label to count

    Raises:
        ValueError: If connection type is not supported
    """
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection

    if isinstance(conn, KuzuConnection):
        from graph.exploration.queries.kuzu.total_counts import (
            get_node_counts_by_group as get_node_counts_by_group_kuzu,
        )
        return get_node_counts_by_group_kuzu(conn)
    elif isinstance(conn, Neo4jConnection):
        from graph.exploration.queries.neo4j.total_counts import (
            get_node_counts_by_group as get_node_counts_by_group_neo4j,
        )
        return get_node_counts_by_group_neo4j(conn)
    else:
        raise ValueError(f"Unsupported connection type: {type(conn).__name__}")


def get_edge_counts_by_group(conn: DbConnection) -> GroupCounts:
    """
    Get count of relationships grouped by type.

    Args:
        conn: Database connection

    Returns:
        Dictionary mapping relationship type to count

    Raises:
        ValueError: If connection type is not supported
    """
    from graph.connections.kuzu.db_connection import KuzuConnection
    from graph.connections.neo4j.db_connection import Neo4jConnection

    if isinstance(conn, KuzuConnection):
        from graph.exploration.queries.kuzu.total_counts import (
            get_edge_counts_by_group as get_edge_counts_by_group_kuzu,
        )
        return get_edge_counts_by_group_kuzu(conn)
    elif isinstance(conn, Neo4jConnection):
        from graph.exploration.queries.neo4j.total_counts import (
            get_edge_counts_by_group as get_edge_counts_by_group_neo4j,
        )
        return get_edge_counts_by_group_neo4j(conn)
    else:
        raise ValueError(f"Unsupported connection type: {type(conn).__name__}")
