"""Bulk export of nodes and edges from AGE graphs to DataFrames."""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import pandas as pd

    from piggie.driver import PiggieDriver


def export_nodes(
    driver: PiggieDriver,
    graph: str,
    label: str,
    *,
    properties: list[str] | None = None,
    limit: int | None = None,
) -> pd.DataFrame:
    """Export nodes of a given label to a pandas DataFrame.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        label: Vertex label to export.
        properties: Specific properties to include. Defaults to all.
        limit: Maximum number of nodes to export.

    Returns:
        DataFrame with expanded node properties.
    """
    query = f"MATCH (n:{label}) RETURN n"
    if limit is not None:
        query += f" LIMIT {limit}"

    result = driver.cypher(query, graph=graph)
    df = result.to_df(expand=True)

    if df.empty:
        return df

    # Filter to requested properties if specified
    if properties is not None:
        keep = []
        for col in df.columns:
            # Keep metadata columns (_id, _label) and requested property columns
            base = col.split(".")[-1] if "." in col else col
            if base.startswith("_") or base in properties:
                keep.append(col)
        df = df[keep]

    return df


def export_edges(
    driver: PiggieDriver,
    graph: str,
    rel_type: str,
    *,
    properties: list[str] | None = None,
    limit: int | None = None,
) -> pd.DataFrame:
    """Export edges of a given type to a pandas DataFrame.

    Args:
        driver: Connected PiggieDriver.
        graph: Graph name.
        rel_type: Edge label to export.
        properties: Specific properties to include. Defaults to all.
        limit: Maximum number of edges to export.

    Returns:
        DataFrame with expanded edge properties.
    """
    query = f"MATCH ()-[r:{rel_type}]->() RETURN r"
    if limit is not None:
        query += f" LIMIT {limit}"

    result = driver.cypher(query, graph=graph)
    df = result.to_df(expand=True)

    if df.empty:
        return df

    if properties is not None:
        keep = []
        for col in df.columns:
            base = col.split(".")[-1] if "." in col else col
            if base.startswith("_") or base in properties:
                keep.append(col)
        df = df[keep]

    return df
