from __future__ import annotations

from collections.abc import Generator
from typing import Any

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.schema.kuzu.schema_extractor import get_node_group_pk_map as get_kuzu_node_group_pk_map


def _extract_node_id(node: object, node_groups_info: dict[str, str]) -> Any:
    if not isinstance(node, dict):
        return pd.NA

    node_label = node.get("_label")
    if not isinstance(node_label, str):
        return pd.NA

    primary_key = node_groups_info.get(node_label)
    if primary_key is None:
        return pd.NA

    return node.get(primary_key, pd.NA)


def get_node_groups_info(conn: DbConnection) -> dict[str, str]:
    return get_kuzu_node_group_pk_map(conn)


def stream_node_metric(
    conn: DbConnection,
    *,
    query_builder: Any,
    value_columns: dict[str, str],
    node_groups_info: dict[str, str] | None = None,
    batch_size: int = 10000,
) -> Generator[pd.DataFrame, Any, None]:
    node_groups_info = node_groups_info or get_node_groups_info(conn)
    offset = 0
    stop = False

    while not stop:
        query = query_builder(offset, batch_size)
        result = conn.execute(query)
        try:
            df = result.get_as_df()
            if df.empty:
                stop = True
                continue

            df["node_group"] = df["node"].apply(lambda x: x.get("_label") if isinstance(x, dict) else pd.NA)
            df["node_id"] = df["node"].apply(lambda x: _extract_node_id(x, node_groups_info))
            df.rename(columns=value_columns, inplace=True)
            columns_to_keep = [*value_columns.values(), "node_group", "node_id"]
            yield df[columns_to_keep]

            if len(df) < batch_size:
                stop = True
            offset += batch_size
        finally:
            result.close()
