from typing import Dict, TypedDict

import kuzu
from pandas import DataFrame
from typing_extensions import TypeAlias

GroupCounts: TypeAlias = Dict[str, int]

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


def get_node_counts_by_group(conn: kuzu.Connection) -> GroupCounts:
    query_result = conn.execute(
        f"""
        MATCH (n)
        RETURN COUNT(n) AS count, LABEL(n) AS node_group
        """
    )
    df: DataFrame = query_result.get_as_df()  # type: ignore
    return {row["node_group"]: row["count"] for _, row in df.iterrows()}

def get_edge_counts_by_group(conn: kuzu.Connection) -> GroupCounts:
    query_result = conn.execute(
        f"""
        MATCH ()-[r]->()
        RETURN COUNT(r) AS count, LABEL(r) AS edge_group
        """
    )
    df: DataFrame = query_result.get_as_df()  # type: ignore
    return {row["edge_group"]: row["count"] for _, row in df.iterrows()}
