from __future__ import annotations

import logging
import os
import tempfile
from functools import cached_property
from pathlib import Path
from typing import Dict, List

import dataiku
from dataiku.core import workload_local_folder
from dataiku.customwebapp import get_webapp_config
from dataikuapi.dss.llm import DSSLLM
from dataikuapi.dss.project import DSSProject

from backends.shared.log_sanitization import sanitize_config_for_logs

logger = logging.getLogger(__name__)

GRAPH_INSTANCES_DIR_NAME = "graph-instances"
WORKLOAD_KUZU_DIR_NAME = "kuzu"
DEBUG_RUN_FOLDER_ENV_VAR_NAME = "DEBUG_RUN_FOLDER"
TMP_EDITOR_KUZU_DIR_PREFIX = "dss-plugin-graph-editor-editor-kuzu-"


class WebAppConfig:
    def __init__(self) -> None:
        self.__is_setup = False
        self.__webapp_config: Dict = {}
        self.__default_project_key = ""
        self.__db_local_folder_path = ""
        self.__db_query_timeout_seconds = 60
        self.__nodes_datasets: List[str] = []
        self.__edges_datasets: List[str] = []
        self.__metadata_ds: str | None = None
        self.__snapshots_ds: str | None = None
        self.__build_graph_recipe_output_connection: str | None = None
        self.__llm_id: str | None = None
        self.__logging_ds: str | None = None
        self.__neo4j_configuration: Dict | None = None

    def setup(self, webapp_config: Dict | None = None, default_project_key: str | None = None) -> None:
        if self.__is_setup:
            return

        self.__webapp_config = webapp_config if webapp_config else get_webapp_config()
        assert self.__webapp_config

        self.__default_project_key = (
            default_project_key if default_project_key else dataiku.get_custom_variables()["projectKey"]
        )
        assert self.__default_project_key

        if logger.isEnabledFor(logging.INFO):
            logger.info("Webapp config is %s.", sanitize_config_for_logs(self.__webapp_config))

        # When debugging locally, use a path passed as an environment variable instead of the scoped temp root.
        debug_run_folder_path = os.getenv(DEBUG_RUN_FOLDER_ENV_VAR_NAME)

        if debug_run_folder_path:
            self.__db_local_folder_path = self.__create_db_folder__(debug_run_folder_path)
        else:
            self.__db_local_folder_path = self.__create_runtime_db_folder__()

        logger.info(f"Kuzu files are persisted locally at '{self.__db_local_folder_path}'.")

        if self.__webapp_config.get("db_query_timeout_seconds"):
            self.__db_query_timeout_seconds = int(self.__webapp_config.get("db_query_timeout_seconds"))  # type: ignore

        if self.__webapp_config.get("nodes_datasets"):
            self.__nodes_datasets = self.__webapp_config.get("nodes_datasets")  # type: ignore

        if self.__webapp_config.get("edges_datasets"):
            self.__edges_datasets = self.__webapp_config.get("edges_datasets")  # type: ignore

        if self.__webapp_config.get("metadata_ds", None):
            self.__metadata_ds = self.__webapp_config.get("metadata_ds", None)  # type: ignore

        if self.__webapp_config.get("snapshots_ds"):
            self.__snapshots_ds = self.__webapp_config.get("snapshots_ds")

        if self.__webapp_config.get("build_graph_recipe_output_connection"):
            self.__build_graph_recipe_output_connection = self.__webapp_config.get(
                "build_graph_recipe_output_connection"
            )

        if self.__webapp_config.get("llm_id"):
            self.__llm_id = self.__webapp_config.get("llm_id")

        if self.__webapp_config.get("logging_ds"):
            self.__logging_ds = self.__webapp_config.get("logging_ds")

        if self.__webapp_config.get("neo4j_configuration"):
            self.__neo4j_configuration = self.__webapp_config.get("neo4j_configuration")

        self.__is_setup = True

    @property
    def default_project_key(self):
        return self.__default_project_key

    @cached_property
    def default_project(self) -> DSSProject:
        client = dataiku.api_client()
        return client.get_project(self.default_project_key)

    @property
    def db_folder_path(self) -> str:
        return self.__db_local_folder_path

    @property
    def db_query_timeout_seconds(self) -> int:
        return self.__db_query_timeout_seconds

    @property
    def nodes_datasets(self) -> List[str]:
        return self.__nodes_datasets

    @property
    def edges_datasets(self) -> List[str]:
        return self.__edges_datasets

    @property
    def metadata_ds(self) -> str | None:
        return self.__metadata_ds

    @property
    def snapshots_ds(self) -> str | None:
        return self.__snapshots_ds

    @property
    def build_graph_recipe_output_connection(self) -> str | None:
        return self.__build_graph_recipe_output_connection

    @property
    def llm_id(self) -> str | None:
        return self.__llm_id

    @property
    def logging_ds(self) -> str | None:
        return self.__logging_ds

    @property
    def neo4j_configuration(self) -> Dict | None:
        return self.__neo4j_configuration

    def get_columns(self, dataset_name: str):
        dataset = dataiku.Dataset(name=dataset_name, project_key=self.default_project_key)
        return dataset.get_config().get("schema").get("columns")

    def __create_db_folder__(self, parent_path: str | Path) -> str:
        p = Path(parent_path) / GRAPH_INSTANCES_DIR_NAME
        p.mkdir(parents=True, exist_ok=True)
        return str(p)

    def __create_runtime_db_folder__(self) -> str:
        if dataiku.is_container_exec():
            logger.info("Containerized execution detected, using an isolated temporary Kuzu path.")
            return self.__create_isolated_tmp_db_folder__()

        return self.__create_db_folder__(self.__get_workload_db_parent())

    def __get_workload_db_parent(self) -> Path:
        return Path(workload_local_folder.get_workload_local_folder_path()) / WORKLOAD_KUZU_DIR_NAME

    def __create_isolated_tmp_db_folder__(self) -> str:
        # In contexts where DSS cannot provide a persistent workload folder (for example containerized
        # execution), keep the temp root isolated per backend start to avoid shared-parent permission issues.
        return self.__create_db_folder__(Path(tempfile.mkdtemp(prefix=TMP_EDITOR_KUZU_DIR_PREFIX)))

    def get_llm(self) -> DSSLLM | None:
        if not self.llm_id:
            return None

        client = dataiku.api_client()
        project = client.get_project(self.default_project_key)
        return project.get_llm(self.llm_id)


webapp_config = WebAppConfig()
