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

import logging

from dataiku import Dataset, Folder
from dataiku.core.dataset_write import DatasetWriter, FakeDatasetWriter
from dataiku.customrecipe import get_input_names_for_role, get_output_names_for_role, get_recipe_config

from backends.explorer.connections.db_config_scanner import DbConfigScanner
from backends.explorer.connections.kuzu.db_handler_factory import KuzuDbHandlerFactory
from backends.explorer.connections.neo4j.db_handler_factory import Neo4jDbHandlerFactory
from graph.connections.neo4j.preset import Neo4jPresetConfig
from graph.exploration.algorithms.page_rank import PageRankParams, ProjectedGraph, compute_page_rank

logger = logging.getLogger(__name__)


db_folder_name = get_input_names_for_role("graph_db_folder")[0]
db_folder = Folder(db_folder_name)
output_ds_name = get_output_names_for_role("output_ds")[0]
output_ds = Dataset(output_ds_name)

config = get_recipe_config()
neo4j_configuration = config.get("neo4j_preset")
if neo4j_configuration:
    neo4j_configuration = Neo4jPresetConfig(**neo4j_configuration)


db_config_scanner = DbConfigScanner([db_folder])
graph_id_list = db_config_scanner.get_snapshot_ids()
if not graph_id_list:
    raise ValueError("No graph found in folder.")
if len(graph_id_list) > 1:
    raise ValueError("More than one graph found in folder.")
graph_id = graph_id_list[0]
db_folder = db_config_scanner.get_db_folder_by_snapshot_id(graph_id)
db_type = db_config_scanner.get_db_type(graph_id)
db_name = db_config_scanner.get_db_name(graph_id)


if db_type == "kuzu":
    db_handler = KuzuDbHandlerFactory(db_folder, graph_id).get_db_handler()

elif db_type == "neo4j":
    if not neo4j_configuration:
        raise ValueError("No Neo4j preset configured.")
    db_handler = Neo4jDbHandlerFactory(db_folder, graph_id, neo4j_configuration).get_db_handler()
else:
    raise ValueError(f"Unknown database type: '{db_type}'")

try:
    with db_handler.get_new_db_instance(readonly=False) as db_instance:
        logger.info(f"Executing compute rank on DB '{db_name}'.")
        with db_instance.get_new_conn() as connection:
            projected_graph = ProjectedGraph("page_rank_graph", config["node_groups"], 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=bool(config.get("normalize_initial", True)),
            )

            logger.info(f"Computing PageRank on DB '{db_name}'.")
            output_ds = Dataset(get_output_names_for_role("output_ds")[0])

            first_batch = True
            output_writer: DatasetWriter | FakeDatasetWriter | None = (
                None  # Initilize once first batch starts, as there are checks that are done beforehand. If these checks raise an exception, the final error is the dataset writer closing which is misleading.
            )
            try:
                for df in compute_page_rank(connection, projected_graph, params, config.get("batch_size", 10000)):
                    if not output_writer:
                        output_writer = output_ds.get_writer()

                    if first_batch:
                        output_ds.write_schema_from_dataframe(df, drop_and_create=True)
                        first_batch = False

                    result_length = len(df)

                    logger.info(f"Writing batch of {result_length} rows to output dataset {output_ds.name}.")
                    output_writer.write_dataframe(df)
                    logger.debug(f"Done writing batch of {result_length} rows to output dataset {output_ds.name}.")
            finally:
                if output_writer:
                    output_writer.__exit__(None, None, None)

        logger.info("Done computing PageRank.")
finally:
    db_handler.cleanup_handler()

