# -*- coding: utf-8 -*-
import logging
import os
from pathlib import Path

import kuzu
from dataiku import Dataset, Folder
from dataiku.core.base import is_container_exec
from dataiku.customrecipe import get_input_names_for_role, get_output_names_for_role, get_recipe_config

from solutions.graph.graph_db_instance_manager import (
    AbstractDbInstance,
    LocalDbInstance,
    LocalReplicaDbInstance,
)

logger = logging.getLogger(__name__)

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

db_instance: AbstractDbInstance

relative_path_to_db = get_recipe_config()["path_to_db"]
db_folder_type = db_folder.get_info(sensitive_info=True)["type"]

if db_folder_type == "Filesystem" and not is_container_exec():
    path_to_db = os.path.join(db_folder.get_path(), relative_path_to_db)
    db_instance = LocalDbInstance(path_to_db, readonly=False)
else:
    path_to_db = Path(relative_path_to_db)
    db_instance = LocalReplicaDbInstance(path_to_db, db_folder, readonly=False)


with db_instance:
    cypher_query = get_recipe_config()["cypher_query"]

    logger.info(f"Executing query on DB in folder {db_folder_name} at {relative_path_to_db}.")

    with db_instance.get_new_conn() as conn_context_manager:
        result = conn_context_manager.connection.execute(cypher_query)
        if isinstance(result, kuzu.QueryResult):
            output_df = result.get_as_df()
            result.close()
        else:
            raise Exception("Multi QueryResult are not supported.")

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

        output_ds.write_with_schema(output_df, drop_and_create=True)

        logger.info("Done executing query.")
