# -*- 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.page_rank import PageRankParams, ProjectedGraph
from recipes.dataiku_recipe_helpers import (
    GraphAnalyticsRecipeBackendConfig,
    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_name = get_output_names_for_role("output_ds")[0]
output_ds = Dataset(output_ds_name)

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

projected_graph = ProjectedGraph(set(config["node_groups"]), set(config["edge_groups"]))
params = PageRankParams(
    damping_factor=float(config.get("damping_factor", 0.85)),
    max_iterations=int(config.get("max_iterations", 100)),
    tolerance=float(config.get("tolerance", 1e-7)),
    normalizeInitial=parse_recipe_bool(config.get("normalize_initial"), default=True),
)

effective_database_type, conn_info, db_handler = backend_config.resolve_backend(db_folder=db_folder)
validate_graph_analytics_recipe_backend(
    "page_rank",
    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(
        "page_rank",
        conn_info,
        db_handler,
        projected_graph,
        params,
        readonly=backend_config.location != GraphAnalyticsExecutionLocation.NATIVE,
        batch_size=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(f"Computing PageRank on DB '{execution.connection_info.uid}'.")

        def log_batch_write(df) -> None:
            result_length = len(df)
            logger.info(f"Writing batch of {result_length} rows to output dataset {output_ds.name}.")

        write_batches_to_dataset(
            output_ds,
            execution.batches,
            on_batch_write=log_batch_write,
        )
except RuntimeError as error:
    raise rewrite_kuzu_recipe_runtime_error(error, effective_database_type) from error

logger.info("Done computing PageRank.")
