from __future__ import annotations

import os

from dataiku.core import workload_local_folder

from backends.editor.utils.webapp_config import webapp_config
from stores.graph_metadata_store import AbstractGraphMetadataStore, DataikuGraphMetadataStore, GraphMetadataStore
from stores.json_file_graph_metadata_store import JsonFileGraphMetadataStore

from .diskcache import get_or_create_diskcache_dir

ENV_VAR_GRAPH_DISK_CACHE_DIR = "VG_GRAPH_DISKCACHE_DIR"
ENV_VAR_SNAPSHOT_DISK_CACHE_DIR = "VG_SNAPSHOT_DISKCACHE_DIR"
ENV_VAR_KUZU_DB_INSTANCE_ACCESS_DISK_CACHE_DIR = "VG_KUZU_DB_INSTANCE_ACCESS_DISKCACHE_DIR"
ENV_VAR_NEO4J_DB_INSTANCE_ACCESS_DISK_CACHE_DIR = "VG_NEO4J_DB_INSTANCE_ACCESS_DISKCACHE_DIR"


def get_or_create_graph_diskcache_dir() -> str:
    """
    Create a shared temporary cache path per webapp run.
    The generated path is shared across processes via an env variable.
    """
    return get_or_create_diskcache_dir(ENV_VAR_GRAPH_DISK_CACHE_DIR, "vg_graph_cache_")


def get_or_create_snapshot_diskcache_dir() -> str:
    """
    Create a shared temporary cache path per webapp run.
    The generated path is shared across processes via an env variable.
    """
    return get_or_create_diskcache_dir(ENV_VAR_SNAPSHOT_DISK_CACHE_DIR, "vg_snapshot_cache_")


def get_or_create_kuzu_db_instance_access_diskcache_dir() -> str:
    """
    Create a shared temporary cache path per webapp run.
    The generated path is shared across processes via an env variable.
    """
    return get_or_create_diskcache_dir(ENV_VAR_KUZU_DB_INSTANCE_ACCESS_DISK_CACHE_DIR, "vg_kuzu_db_instance_access_")


def get_or_create_neo4j_db_instance_access_diskcache_dir() -> str:
    """
    Create a shared temporary cache path per webapp run.
    The generated path is shared across processes via an env variable.
    """
    return get_or_create_diskcache_dir(ENV_VAR_NEO4J_DB_INSTANCE_ACCESS_DISK_CACHE_DIR, "vg_neo4j_db_instance_access_")


def create_graph_metadata_store() -> AbstractGraphMetadataStore:
    """
    Graph metadata dataset is optional. If not configured, use the workload folder of the webapp to persist graph metadata.
    """
    diskcache_path = get_or_create_graph_diskcache_dir()
    if metadata_ds := webapp_config.metadata_ds:
        return GraphMetadataStore(
            diskcache_path,
            DataikuGraphMetadataStore(metadata_ds, webapp_config.default_project_key),
        )
    else:
        storage_dir = os.path.join(workload_local_folder.get_workload_local_folder_path(), "graphs_metadata")
        return GraphMetadataStore(
            diskcache_path,
            JsonFileGraphMetadataStore(storage_dir),
        )
