# -*- 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.connections.neo4j.preset import Neo4jPresetConfig
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__)

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.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

db_folder_names = get_input_names_for_role("graph_db_folder")
db_folder = Folder(db_folder_names[0]) if db_folder_names else 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=True) as db_instance:
        db_display_name = getattr(conn_info, "database_name", conn_info.uid)
        logger.info(f"Executing query on DB '{db_display_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()
except RuntimeError as error:
    raise rewrite_kuzu_recipe_runtime_error(error, effective_database_type) from error
finally:
    db_handler.cleanup_handler()
