from functools import wraps

from editor.backend.utils.store import get_or_create_graph_diskcache_dir, get_or_create_snapshot_diskcache_dir
from editor.backend.utils.webapp_config import webapp_config
from solutions.graph.store.graph_metadata_snapshot_store import (
    DataikuGraphMetadataSnapshotStore,
    GraphMetadataSnapshotStore,
)
from solutions.graph.store.graph_metadata_store import DataikuGraphMetadataStore, GraphMetadataStore


def use_graph_metadata_store(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        diskcache_path = get_or_create_graph_diskcache_dir()
        graph_store = GraphMetadataStore(
            diskcache_path,
            DataikuGraphMetadataStore(webapp_config.metadata_ds, webapp_config.default_project_key),
        )

        kwargs["graph_store"] = graph_store

        return f(*args, **kwargs)

    return wrapper


def use_graph_snapshot_store(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if webapp_config.snapshots_ds:
            diskcache_path = get_or_create_snapshot_diskcache_dir()
            snapshot_store = GraphMetadataSnapshotStore(
                diskcache_path,
                DataikuGraphMetadataSnapshotStore(webapp_config.snapshots_ds, webapp_config.default_project_key),
            )
        else:
            snapshot_store = None

        kwargs["snapshot_store"] = snapshot_store

        return f(*args, **kwargs)

    return wrapper
