from __future__ import annotations

from typing import Any, Generator

import pandas as pd

from graph.connections.db_connection import DbConnection
from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsExecutionLocation,
    GraphAnalyticsSamplingMode,
    get_graph_analytics_backend_service,
)
from graph.exploration.algorithms.graph_features_params import GraphFeaturesParams, selected_graph_feature_names
from graph.exploration.algorithms.in_memory.graph_features.common import build_empty_output_frame
from graph.exploration.algorithms.projected_graph import ProjectedGraph


def compute_graph_features(
    conn: DbConnection,
    projected_graph: ProjectedGraph,
    params: GraphFeaturesParams,
    batch_size: int = 10000,
    *,
    location: GraphAnalyticsExecutionLocation = GraphAnalyticsExecutionLocation.IN_MEMORY,
    sampling_mode: GraphAnalyticsSamplingMode | None = None,
    sampling_limit: int | None = None,
) -> Generator[pd.DataFrame, Any, None]:
    return get_graph_analytics_backend_service().execute(
        "graph_features",
        conn,
        projected_graph,
        params,
        batch_size=batch_size,
        location=location,
        sampling_mode=sampling_mode,
        sampling_limit=sampling_limit,
    )


def build_empty_graph_features_output_frame(params: GraphFeaturesParams) -> pd.DataFrame:
    return build_empty_output_frame(params)


__all__ = [
    "GraphFeaturesParams",
    "ProjectedGraph",
    "build_empty_graph_features_output_frame",
    "compute_graph_features",
    "selected_graph_feature_names",
]
