from functools import wraps
from typing import Callable, Optional

from backends.editor.connections.db_handler_registry import get_handler_for_graph
from backends.editor.utils.store import (
    create_graph_metadata_store,
    get_or_create_graph_diskcache_dir,
    get_or_create_snapshot_diskcache_dir,
)
from backends.editor.utils.webapp_config import webapp_config
from solutions.graph.models import GraphId
from stores.graph_metadata_snapshot_store import (
    DataikuGraphMetadataSnapshotStore,
    GraphMetadataSnapshotStore,
)


def use_graph_metadata_store(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        diskcache_path = get_or_create_graph_diskcache_dir()
        graph_store = create_graph_metadata_store()

        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


def use_db_handler(id_getter: Callable[[], Optional[GraphId]]):
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            graph_id = id_getter()

            if not graph_id:
                raise ValueError("Could not resolve graph_id from request")

            # Get the handler from the registry (cached via LRU)
            db_handler = get_handler_for_graph(graph_id)

            kwargs["db_handler"] = db_handler

            return f(*args, **kwargs)

        return wrapper

    return decorator
