from __future__ import annotations

from pathlib import Path
from typing import Dict

from llama_index.core import StorageContext, VectorStoreIndex
from llama_index.core.graph_stores.simple import SimpleGraphStore
from llama_index.core.graph_stores.simple_labelled import SimplePropertyGraphStore
from llama_index.core.graph_stores.types import GraphStore, PropertyGraphStore
from llama_index.core.storage.docstore.simple_docstore import SimpleDocumentStore
from llama_index.core.storage.docstore.types import BaseDocumentStore
from llama_index.core.storage.index_store.simple_index_store import SimpleIndexStore
from llama_index.core.storage.index_store.types import BaseIndexStore
from llama_index.core.vector_stores.simple import SimpleVectorStore
from llama_index.core.vector_stores.types import BasePydanticVectorStore


class LocalStorageBackend:
    """Filesystem-backed StorageContext factory + persistence hooks."""

    def __init__(self) -> None:
        self._storage_contexts: Dict[str, StorageContext] = {}

    def _get_vector_store(self, persist_dir: Path) -> BasePydanticVectorStore:
        return SimpleVectorStore(persist_dir=str(persist_dir))

    def _get_index_store(self, persist_dir: Path) -> BaseIndexStore:
        try:
            return SimpleIndexStore.from_persist_dir(persist_dir=str(persist_dir))
        except Exception:
            return SimpleIndexStore()

    def _get_docstore(self, persist_dir: Path) -> BaseDocumentStore:
        try:
            return SimpleDocumentStore.from_persist_dir(persist_dir=str(persist_dir))
        except Exception:
            return SimpleDocumentStore()

    def _get_graph_store(self, persist_dir: Path) -> GraphStore:
        try:
            return SimpleGraphStore.from_persist_dir(persist_dir=str(persist_dir))
        except Exception:
            return SimpleGraphStore()

    def _get_property_graph_store(self, persist_dir: Path) -> PropertyGraphStore:
        try:
            return SimplePropertyGraphStore.from_persist_dir(persist_dir=str(persist_dir))
        except Exception:
            return SimplePropertyGraphStore()

    def create_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> StorageContext:
        scoped_dir = index_dir / name
        scoped_dir.mkdir(parents=True, exist_ok=True)
        return StorageContext.from_defaults(
            vector_store=self._get_vector_store(persist_dir=scoped_dir),
            index_store=self._get_index_store(persist_dir=scoped_dir),
            docstore=self._get_docstore(persist_dir=data_dir),
            graph_store=self._get_graph_store(persist_dir=scoped_dir),
            property_graph_store=self._get_property_graph_store(persist_dir=scoped_dir),
            persist_dir=str(scoped_dir),
        )

    def get_storage_context(self, *, index_dir: Path, data_dir: Path, name: str) -> StorageContext:
        if name in self._storage_contexts:
            return self._storage_contexts[name]
        storage_context = self.create_storage_context(index_dir=index_dir, data_dir=data_dir, name=name)
        self._storage_contexts[name] = storage_context
        return storage_context

    def save(self, *, vector_index: VectorStoreIndex, index_dir: Path, data_dir: Path) -> None:
        vector_index.storage_context.docstore.persist(persist_path=str(data_dir / "docstore.json"))
        vector_index.storage_context.persist(persist_dir=index_dir)

