# -*- 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 graph.connections.neo4j.preset import Neo4jPresetConfig
from graph.exploration.algorithms.page_rank import PageRankParams, ProjectedGraph, compute_page_rank
from recipes.explore_graph.utils import get_db_handler, get_single_connection_info
from recipes.runtime_errors import rewrite_kuzu_recipe_runtime_error
from solutions.graph.models import DbType

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()
neo4j_configuration = config.get("neo4j_preset")
if neo4j_configuration:
    neo4j_configuration = Neo4jPresetConfig.from_dict(neo4j_configuration)

database_name = config.get("database_name")
database_type = config.get("database_type")

effective_database_type: DbType | None
if database_type == "unmanaged_neo4j":
    effective_database_type = "neo4j"
elif database_type in ("kuzu", "neo4j"):
    effective_database_type = database_type
else:
    effective_database_type = None

conn_info = get_single_connection_info(
    db_folder=db_folder,
    neo4j_configuration=neo4j_configuration,
    database_name=database_name,
    database_type=effective_database_type,
)
db_handler = get_db_handler(
    db_folder=db_folder,
    neo4j_configuration=neo4j_configuration,
    database_name=database_name,
    database_type=effective_database_type,
    snapshot_id=conn_info.uid,
)

try:
    with db_handler.get_new_db_instance(readonly=False) as db_instance:
        logger.info(f"Executing compute rank on DB '{conn_info.uid}'.")
        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 '{conn_info.uid}'.")
            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.")
except RuntimeError as error:
    raise rewrite_kuzu_recipe_runtime_error(error, effective_database_type) from error
finally:
    db_handler.cleanup_handler()
