# -*- coding: utf-8 -*-
from __future__ import annotations

import logging

from dataiku import Dataset, Folder
from dataiku.customrecipe import get_input_names_for_role, get_output_names_for_role, get_recipe_config

from graph.exploration.algorithms.backend_service import (
    GraphAnalyticsExecutionLocation,
    build_recipe_execution_note,
    execute_graph_algorithm,
)
from graph.exploration.algorithms.graph_algorithm_registry import (
    GRAPH_FEATURE_ALGORITHM_DEFINITIONS,
)
from graph.exploration.algorithms.graph_features import (
    GraphFeaturesParams,
    ProjectedGraph,
    build_empty_graph_features_output_frame,
)
from recipes.dataiku_recipe_helpers import (
    GraphAnalyticsRecipeBackendConfig,
    parse_graph_algorithm_selection,
    parse_recipe_bool,
    validate_graph_analytics_recipe_backend,
    write_batches_to_dataset,
)
from recipes.runtime_errors import rewrite_kuzu_recipe_runtime_error

logger = logging.getLogger(__name__)


db_folder_names = get_input_names_for_role("graph_db_folder")
db_folder = Folder(db_folder_names[0]) if db_folder_names else None
output_ds = Dataset(get_output_names_for_role("output_ds")[0])

config = get_recipe_config()
backend_config = GraphAnalyticsRecipeBackendConfig.from_recipe_config(
    config,
    default_location=GraphAnalyticsExecutionLocation.IN_MEMORY,
)

is_directed = parse_recipe_bool(config.get("directed_graph"), default=False)
selected_algorithms = set(
    parse_graph_algorithm_selection(
        config,
        definitions=GRAPH_FEATURE_ALGORITHM_DEFINITIONS,
        directed_graph=is_directed,
    )
)

params = GraphFeaturesParams(
    directed_graph=is_directed,
    output_type=config.get("output_type", "output_nodes"),
    degree="degree" in selected_algorithms,
    eigenvector_centrality="eigenvector_centrality" in selected_algorithms,
    clustering="clustering" in selected_algorithms,
    triangles="triangles" in selected_algorithms,
    closeness="closeness" in selected_algorithms,
    pagerank="pagerank" in selected_algorithms,
    sq_clustering="sq_clustering" in selected_algorithms,
    connected_components="connected_components" in selected_algorithms,
)

projected_graph = ProjectedGraph(set(config.get("node_groups", [])), set(config.get("edge_groups", [])))

logger.info(
    "Computing graph features on selected groups (directed=%s, output_type=%s).",
    params.directed_graph,
    params.output_type,
)

effective_database_type, conn_info, db_handler = backend_config.resolve_backend(db_folder=db_folder)
validate_graph_analytics_recipe_backend(
    "graph_features",
    effective_database_type=effective_database_type,
    backend_config=backend_config,
)

try:
    logger.info(build_recipe_execution_note(backend_config.location, backend_config.sampling_mode))
    with execute_graph_algorithm(
        "graph_features",
        conn_info,
        db_handler,
        projected_graph,
        params,
        readonly=True,
        batch_size=int(config.get("batch_size", 10000)),
        location=backend_config.location,
        sampling_mode=backend_config.sampling_mode,
        sampling_limit=backend_config.sampling_limit,
    ) as execution:
        logger.info("Computing graph features on DB '%s'.", execution.connection_info.uid)
        write_batches_to_dataset(
            output_ds,
            execution.batches,
            empty_frame_factory=lambda: build_empty_graph_features_output_frame(params),
        )
except RuntimeError as error:
    raise rewrite_kuzu_recipe_runtime_error(error, effective_database_type) from error

logger.info("Done computing graph features.")
