# -*- coding: utf-8 -*-
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 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

logger = logging.getLogger(__name__)

output_ds_name = get_output_names_for_role("output_ds")[0]
output_ds = Dataset(output_ds_name)

config = get_recipe_config()

cypher_query = config.get("cypher_query")
if not cypher_query:
    raise ValueError("Cypher query is not specified.")

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

database_name = config.get("database_name")

db_folder_name = get_input_names_for_role("graph_db_folder")[0]
db_folder = Folder(db_folder_name)

db_config_scanner = DbConfigScanner([db_folder])

graph_id_list = db_config_scanner.get_snapshot_ids()
if not graph_id_list:
    raise ValueError("No graph database found in the input folder.")
if len(graph_id_list) > 1:
    raise ValueError("Multiple graph databases found in the input folder, please select a folder with a single graph.")

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)
if db_type == "kuzu":
    db_handler = KuzuDbHandlerFactory(db_folder, graph_id).get_db_handler()
elif db_type == "neo4j":
    if not neo4j_configuration:
        raise ValueError(f"No Neo4j configuration found for recipe")

    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=True) as db_instance:
        logger.info(f"Executing query on DB '{database_name}'.")
        with db_instance.get_new_conn() as connection:
            result = connection.execute(cypher_query)
            try:
                output_df = result.get_as_df()
                output_ds.write_with_schema(output_df, drop_and_create=True)
                logger.info("Successfully executed query and wrote to output dataset.")
            finally:
                result.close()
finally:
    db_handler.cleanup_handler()
