import os
import re
import json
import time
from dataclasses import dataclass, field
import numpy as np
import configparser
import asyncio

from typing import Any, Union, final

from ..base import (
    BaseGraphStorage,
    BaseKVStorage,
    BaseVectorStorage,
    DocProcessingStatus,
    DocStatus,
    DocStatusStorage,
)
from ..utils import (
    logger,
    compute_mdhash_id,
    _cooperative_yield,
    merge_source_ids,
    validate_workspace,
)
from ..types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
from ..constants import GRAPH_FIELD_SEP, DEFAULT_QUERY_PRIORITY
from .._version import __version__
from ..kg.shared_storage import get_data_init_lock, get_namespace_lock

import pipmaster as pm

if not pm.is_installed("pymongo"):
    pm.install("pymongo")

from pymongo import AsyncMongoClient  # type: ignore
from pymongo import UpdateOne  # type: ignore
from pymongo.asynchronous.database import AsyncDatabase  # type: ignore
from pymongo.asynchronous.collection import AsyncCollection  # type: ignore
from pymongo.operations import SearchIndexModel  # type: ignore
from pymongo.driver_info import DriverInfo  # type: ignore
from pymongo.errors import (  # type: ignore
    PyMongoError,
    DuplicateKeyError,
    BulkWriteError,
)

config = configparser.ConfigParser()
config.read("config.ini", "utf-8")

GRAPH_BFS_MODE = os.getenv("MONGO_GRAPH_BFS_MODE", "bidirectional")

# Flush-time batching limits shared by every MongoDB upsert path
# (MongoVectorDBStorage, MongoKVStorage, MongoGraphStorage).
# The payload-byte budget is the primary limiter; the record-count caps are a
# secondary guard that only binds when individual records are small.
# Upsert and delete have separate count caps on purpose: upsert records each
# carry a full embedding vector and are far heavier than delete _ids, so the
# upsert batch count is kept much smaller than the delete one.
# MongoDB caps a single BSON document at 16MB and a single bulk command message
# at 48MB; a 16MB JSON estimate (which overestimates the real BSON size) keeps
# every bulk_write comfortably below the wire limit and bounds peak memory.
DEFAULT_MONGO_UPSERT_MAX_PAYLOAD_BYTES = 16 * 1024 * 1024  # 16MB
DEFAULT_MONGO_UPSERT_MAX_RECORDS_PER_BATCH = 128
DEFAULT_MONGO_DELETE_MAX_RECORDS_PER_BATCH = 1000

# MongoDB duplicate-key error code, raised when an upsert insert races the
# unique edge-endpoint index (another writer inserted the same edge first).
_DUPLICATE_KEY_CODE = 11000

# Emit a migration progress line every this many deduped docs, so operators
# watching a large migration see liveness (mirrors the OpenSearch canonical-id
# migration's progress cadence).
_EDGE_MIGRATION_PROGRESS_INTERVAL = 50_000


def _canonical_edge_endpoints(
    source_node_id: str, target_node_id: str
) -> tuple[str, str]:
    """Direction-independent ``(edge_lo, edge_hi)`` endpoints for an undirected edge.

    The sorted pair maps ``(A,B)`` and ``(B,A)`` to the same two field values,
    so a *compound* unique index on ``(edge_lo, edge_hi)`` lets MongoDB reject
    the second of two racing inserts (the classic ``$or``-upsert duplicate gap)
    regardless of direction. Storing the endpoints as two separate fields — not
    a single delimiter-joined string — avoids any collision between distinct
    pairs whose ids happen to contain the delimiter (e.g. custom-KG ids), and
    needs no input sanitisation. Reads keep using the bidirectional ``$or``.
    """
    return tuple(sorted((source_node_id, target_node_id)))  # type: ignore[return-value]


def _edge_source_id_list(doc: dict[str, Any]) -> list[str]:
    """Return an edge doc's source ids, from the ``source_ids`` array or by
    splitting the ``GRAPH_FIELD_SEP``-joined ``source_id`` string."""
    sids = doc.get("source_ids")
    if not sids and doc.get("source_id"):
        sids = doc["source_id"].split(GRAPH_FIELD_SEP)
    return list(sids or [])


def _coerce_weight(weight: Any) -> float | None:
    """Coerce a (possibly string) edge weight to float, or None if non-numeric."""
    if weight is None:
        return None
    try:
        return float(weight)
    except (TypeError, ValueError):
        return None


def _estimate_doc_bytes(doc: Any) -> int:
    """Estimate a document's serialized byte size via compact JSON.

    JSON overestimates the real BSON size MongoDB writes (a JSON float string is
    far longer than the 8 bytes a BSON double encodes), so callers stay
    conservatively below server limits and never underestimate.

    This is a splitting *heuristic*, not the exact wire size: upsert callers pass
    only the dominant payload field (the ``$set`` body / ``update_doc``), not the
    full ``UpdateOne`` op (filter, ``$setOnInsert``, ``$or`` wrapper). Those extras
    are tiny next to an embedding/document body, and the 16MB estimate budget sits
    far under MongoDB's 48MB bulk-command limit, so the under-count is immaterial;
    the server stays the final arbiter.
    """
    return len(
        json.dumps(doc, ensure_ascii=False, separators=(",", ":"), default=str).encode(
            "utf-8"
        )
    )


def _chunk_by_budget(
    items: list[Any],
    size_of,
    max_payload_bytes: int,
    max_records_per_batch: int,
) -> list[tuple[list[Any], int]]:
    """Split items into batches by estimated payload size (primary) and count.

    The byte budget is the primary limiter: items accumulate until adding the
    next one would exceed ``max_payload_bytes``, then a new batch starts.
    ``size_of(item)`` returns an item's estimated serialized byte size. A single
    item larger than the byte budget is emitted as its own batch rather than
    raising; the server stays the final arbiter. A non-positive limit disables
    that dimension. Returns ``(batch, summed_estimated_bytes)`` tuples (the
    estimate is used for logging).
    """
    if not items:
        return []

    payload_limit = max_payload_bytes if max_payload_bytes > 0 else float("inf")
    records_limit = max_records_per_batch if max_records_per_batch > 0 else float("inf")

    batches: list[tuple[list[Any], int]] = []
    current: list[Any] = []
    # JSON array overhead ("[]")
    current_bytes = 2

    for item in items:
        item_bytes = size_of(item)
        # If current batch not empty, a comma is needed before next element.
        separator_overhead = 1 if current else 0
        next_bytes = current_bytes + separator_overhead + item_bytes

        if current and (len(current) >= records_limit or next_bytes > payload_limit):
            batches.append((current, current_bytes))
            current = []
            current_bytes = 2
            next_bytes = current_bytes + item_bytes

        current.append(item)
        current_bytes = next_bytes

    if current:
        batches.append((current, current_bytes))

    return batches


def _resolve_upsert_batch_limits() -> tuple[int, int]:
    """Resolve flush-time upsert batching limits from env, with module defaults.

    Shared by every MongoDB upsert path so the byte/record caps that bound a
    single ``bulk_write`` are consistent across all of them. A non-positive
    value disables that splitting dimension.
    """
    max_payload_bytes = int(
        os.getenv(
            "MONGO_UPSERT_MAX_PAYLOAD_BYTES",
            str(DEFAULT_MONGO_UPSERT_MAX_PAYLOAD_BYTES),
        )
    )
    max_records_per_batch = int(
        os.getenv(
            "MONGO_UPSERT_MAX_RECORDS_PER_BATCH",
            str(DEFAULT_MONGO_UPSERT_MAX_RECORDS_PER_BATCH),
        )
    )
    if max_payload_bytes <= 0:
        logger.warning(
            f"MONGO_UPSERT_MAX_PAYLOAD_BYTES={max_payload_bytes} is non-positive, disable payload-size splitting"
        )
    if max_records_per_batch <= 0:
        logger.warning(
            f"MONGO_UPSERT_MAX_RECORDS_PER_BATCH={max_records_per_batch} is non-positive, disable upsert record-count splitting"
        )
    return max_payload_bytes, max_records_per_batch


def _resolve_delete_batch_limit() -> int:
    """Resolve the flush-time delete record-count cap from env, with module default.

    Shared by every MongoDB delete path that fans a list of match clauses into a
    single server message (``delete_many`` with ``$in``/``$or``), so the cap that
    keeps one delete under the bulk message / 16MB query limit is consistent. A
    non-positive value disables record-count splitting.
    """
    max_records_per_batch = int(
        os.getenv(
            "MONGO_DELETE_MAX_RECORDS_PER_BATCH",
            str(DEFAULT_MONGO_DELETE_MAX_RECORDS_PER_BATCH),
        )
    )
    if max_records_per_batch <= 0:
        logger.warning(
            f"MONGO_DELETE_MAX_RECORDS_PER_BATCH={max_records_per_batch} is non-positive, disable delete record-count splitting"
        )
    return max_records_per_batch


async def _run_batched_bulk_write(
    collection,
    ops: list[tuple[Any, int, str]],
    *,
    max_payload_bytes: int,
    max_records_per_batch: int,
    ordered: bool,
    log_prefix: str,
    what: str,
) -> None:
    """Execute UpdateOne ops as payload-size/record-count bounded bulk_write batches.

    ``ops`` is a list of ``(operation, estimated_bytes, id_for_log)`` triples.
    Splitting keeps each bulk command below MongoDB's 48MB message ceiling and
    bounds the in-memory op list. Fail-fast: a batch failure raises and no
    further batches run, so callers must treat the whole write as retryable
    (UpdateOne(..., upsert=True) is idempotent).
    """
    if not ops:
        return

    batches = _chunk_by_budget(
        ops, lambda triple: triple[1], max_payload_bytes, max_records_per_batch
    )
    if len(batches) > 1:
        logger.info(
            f"{log_prefix} {what} split into {len(batches)} batches "
            f"for {len(ops)} records"
        )
    for batch_index, (batch, estimated_bytes) in enumerate(batches, 1):
        if (
            len(batch) == 1
            and max_payload_bytes > 0
            and estimated_bytes > max_payload_bytes
        ):
            logger.warning(
                f"{log_prefix} {what}: single record id={batch[0][2]} "
                f"estimated {estimated_bytes} bytes exceeds {max_payload_bytes}"
            )
        logger.debug(
            f"{log_prefix} {what} batch {batch_index}/{len(batches)}: "
            f"records={len(batch)}, estimated_payload_bytes={estimated_bytes}"
        )
        await collection.bulk_write([triple[0] for triple in batch], ordered=ordered)


class ClientManager:
    _instances: dict = {"client": None, "db": None, "ref_count": 0}
    _lock = asyncio.Lock()

    @classmethod
    async def get_client(cls) -> AsyncMongoClient:
        async with cls._lock:
            if cls._instances["db"] is None:
                uri = os.environ.get(
                    "MONGO_URI",
                    config.get(
                        "mongodb",
                        "uri",
                        fallback="mongodb://root:root@localhost:27017/",
                    ),
                )
                database_name = os.environ.get(
                    "MONGO_DATABASE",
                    config.get("mongodb", "database", fallback="LightRAG"),
                )
                client = AsyncMongoClient(
                    uri,
                    driver=DriverInfo(name="LightRAG", version=__version__),
                )
                db = client.get_database(database_name)
                cls._instances["client"] = client
                cls._instances["db"] = db
                cls._instances["ref_count"] = 0
            cls._instances["ref_count"] += 1
            return cls._instances["db"]

    @classmethod
    async def release_client(cls, db: AsyncDatabase):
        async with cls._lock:
            if db is not None:
                if db is cls._instances["db"]:
                    cls._instances["ref_count"] -= 1
                    if cls._instances["ref_count"] == 0:
                        client = cls._instances.get("client")
                        if client is not None:
                            await client.close()
                        cls._instances["client"] = None
                        cls._instances["db"] = None


@final
@dataclass
class MongoKVStorage(BaseKVStorage):
    db: AsyncDatabase = field(default=None)
    _data: AsyncCollection = field(default=None)

    def __init__(self, namespace, global_config, embedding_func, workspace=None):
        super().__init__(
            namespace=namespace,
            workspace=workspace or "",
            global_config=global_config,
            embedding_func=embedding_func,
        )
        self.__post_init__()

    def __post_init__(self):
        validate_workspace(self.workspace)
        # Check for MONGODB_WORKSPACE environment variable first (higher priority)
        # This allows administrators to force a specific workspace for all MongoDB storage instances
        mongodb_workspace = os.environ.get("MONGODB_WORKSPACE")
        if mongodb_workspace and mongodb_workspace.strip():
            # Use environment variable value, overriding the passed workspace parameter
            effective_workspace = mongodb_workspace.strip()
            logger.info(
                f"Using MONGODB_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
            )
        else:
            # Use the workspace parameter passed during initialization
            effective_workspace = self.workspace
            if effective_workspace:
                logger.debug(
                    f"Using passed workspace parameter: '{effective_workspace}'"
                )

        # Build final_namespace with workspace prefix for data isolation
        # Keep original namespace unchanged for type detection logic
        if effective_workspace:
            self.final_namespace = f"{effective_workspace}_{self.namespace}"
            self.workspace = effective_workspace
            logger.debug(
                f"Final namespace with workspace prefix: '{self.final_namespace}'"
            )
        else:
            # When workspace is empty, final_namespace equals original namespace
            self.final_namespace = self.namespace
            self.workspace = ""
            logger.debug(
                f"[{self.workspace}] Final namespace (no workspace): '{self.namespace}'"
            )

        self._collection_name = self.final_namespace
        (
            self._max_upsert_payload_bytes,
            self._max_upsert_records_per_batch,
        ) = _resolve_upsert_batch_limits()

    async def initialize(self):
        async with get_data_init_lock():
            if self.db is None:
                self.db = await ClientManager.get_client()

            self._data = await get_or_create_collection(self.db, self._collection_name)
            logger.debug(
                f"[{self.workspace}] Use MongoDB as KV {self._collection_name}"
            )

    async def finalize(self):
        if self.db is not None:
            await ClientManager.release_client(self.db)
            self.db = None
            self._data = None

    async def get_by_id(self, id: str) -> dict[str, Any] | None:
        # Unified handling for flattened keys
        doc = await self._data.find_one({"_id": id})
        if doc:
            # Ensure time fields are present, provide default values for old data
            doc.setdefault("create_time", 0)
            doc.setdefault("update_time", 0)
        return doc

    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
        cursor = self._data.find({"_id": {"$in": ids}})
        docs = await cursor.to_list(length=None)

        doc_map: dict[str, dict[str, Any]] = {}
        for doc in docs:
            if not doc:
                continue
            doc.setdefault("create_time", 0)
            doc.setdefault("update_time", 0)
            doc_map[str(doc.get("_id"))] = doc

        ordered_results: list[dict[str, Any] | None] = []
        for id_value in ids:
            ordered_results.append(doc_map.get(str(id_value)))
        return ordered_results

    async def filter_keys(self, keys: set[str]) -> set[str]:
        cursor = self._data.find({"_id": {"$in": list(keys)}}, {"_id": 1})
        existing_ids = {str(x["_id"]) async for x in cursor}
        return keys - existing_ids

    async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
        logger.debug(f"[{self.workspace}] Inserting {len(data)} to {self.namespace}")
        if not data:
            return

        # Unified handling for all namespaces with flattened keys. KV docs
        # (full_docs, text_chunks, llm_response_cache) can be large, so the
        # upsert is split into payload-bounded bulk_write batches.
        operations: list[tuple[Any, int, str]] = []
        current_time = int(time.time())  # Get current Unix timestamp

        for i, (k, v) in enumerate(data.items(), start=1):
            # For text_chunks namespace, ensure llm_cache_list field exists
            if self.namespace.endswith("text_chunks"):
                if "llm_cache_list" not in v:
                    v["llm_cache_list"] = []

            # Create a copy of v for $set operation, excluding create_time to avoid conflicts
            v_for_set = v.copy()
            v_for_set["_id"] = k  # Use flattened key as _id
            v_for_set["update_time"] = current_time  # Always update update_time

            # Remove create_time from $set to avoid conflict with $setOnInsert
            v_for_set.pop("create_time", None)

            operations.append(
                (
                    UpdateOne(
                        {"_id": k},
                        {
                            "$set": v_for_set,  # Update all fields except create_time
                            "$setOnInsert": {
                                "create_time": current_time
                            },  # Set create_time only on insert
                        },
                        upsert=True,
                    ),
                    _estimate_doc_bytes(v_for_set),
                    k,
                )
            )
            await _cooperative_yield(i)

        # ordered=False (intentional): the old single bulk_write used pymongo's
        # default ordered=True, but every op targets a distinct flattened _id, so
        # the writes are order-independent. ordered=False lets the server apply
        # them in parallel and is the right choice for idempotent upserts.
        await _run_batched_bulk_write(
            self._data,
            operations,
            max_payload_bytes=self._max_upsert_payload_bytes,
            max_records_per_batch=self._max_upsert_records_per_batch,
            ordered=False,
            log_prefix=f"[{self.workspace}] {self.namespace} upsert:",
            what="upsert",
        )

    async def index_done_callback(self) -> None:
        # Mongo handles persistence automatically
        pass

    async def is_empty(self) -> bool:
        """Check if the storage is empty for the current workspace and namespace

        Returns:
            bool: True if storage is empty, False otherwise
        """
        try:
            # Use count_documents with limit 1 for efficiency
            count = await self._data.count_documents({}, limit=1)
            return count == 0
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error checking if storage is empty: {e}")
            return True

    async def delete(self, ids: list[str]) -> None:
        """Delete documents with specified IDs

        Args:
            ids: List of document IDs to be deleted
        """
        if not ids:
            return

        # Convert to list if it's a set (MongoDB BSON cannot encode sets)
        if isinstance(ids, set):
            ids = list(ids)

        try:
            result = await self._data.delete_many({"_id": {"$in": ids}})
            logger.info(
                f"[{self.workspace}] Deleted {result.deleted_count} documents from {self.namespace}"
            )
        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error deleting documents from {self.namespace}: {e}"
            )

    async def drop(self) -> dict[str, str]:
        """Drop the storage by removing all documents in the collection.

        Returns:
            dict[str, str]: Status of the operation with keys 'status' and 'message'
        """
        try:
            result = await self._data.delete_many({})
            deleted_count = result.deleted_count

            logger.info(
                f"[{self.workspace}] Dropped {deleted_count} documents from doc status {self._collection_name}"
            )
            return {
                "status": "success",
                "message": f"{deleted_count} documents dropped",
            }
        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error dropping doc status {self._collection_name}: {e}"
            )
            return {"status": "error", "message": str(e)}


@final
@dataclass
class MongoDocStatusStorage(DocStatusStorage):
    db: AsyncDatabase = field(default=None)
    _data: AsyncCollection = field(default=None)

    def _prepare_doc_status_data(self, doc: dict[str, Any]) -> dict[str, Any]:
        """Normalize and migrate a raw Mongo document to DocProcessingStatus-compatible dict."""
        # Make a copy of the data to avoid modifying the original
        data = doc.copy()
        # Remove deprecated content field if it exists
        data.pop("content", None)
        # Remove MongoDB _id field if it exists
        data.pop("_id", None)
        # If file_path is not in data, use document id as file path
        if "file_path" not in data:
            data["file_path"] = "no-file-path"
        # Ensure new fields exist with default values
        if "metadata" not in data:
            data["metadata"] = {}
        if "error_msg" not in data:
            data["error_msg"] = None
        # Backward compatibility: migrate legacy 'error' field to 'error_msg'
        if "error" in data:
            if "error_msg" not in data or data["error_msg"] in (None, ""):
                data["error_msg"] = data.pop("error")
            else:
                data.pop("error", None)
        return data

    def __init__(self, namespace, global_config, embedding_func, workspace=None):
        super().__init__(
            namespace=namespace,
            workspace=workspace or "",
            global_config=global_config,
            embedding_func=embedding_func,
        )
        self.__post_init__()

    def __post_init__(self):
        validate_workspace(self.workspace)
        # Check for MONGODB_WORKSPACE environment variable first (higher priority)
        # This allows administrators to force a specific workspace for all MongoDB storage instances
        mongodb_workspace = os.environ.get("MONGODB_WORKSPACE")
        if mongodb_workspace and mongodb_workspace.strip():
            # Use environment variable value, overriding the passed workspace parameter
            effective_workspace = mongodb_workspace.strip()
            logger.info(
                f"Using MONGODB_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
            )
        else:
            # Use the workspace parameter passed during initialization
            effective_workspace = self.workspace
            if effective_workspace:
                logger.debug(
                    f"Using passed workspace parameter: '{effective_workspace}'"
                )

        # Build final_namespace with workspace prefix for data isolation
        # Keep original namespace unchanged for type detection logic
        if effective_workspace:
            self.final_namespace = f"{effective_workspace}_{self.namespace}"
            self.workspace = effective_workspace
            logger.debug(
                f"Final namespace with workspace prefix: '{self.final_namespace}'"
            )
        else:
            # When workspace is empty, final_namespace equals original namespace
            self.final_namespace = self.namespace
            self.workspace = ""
            logger.debug(f"Final namespace (no workspace): '{self.final_namespace}'")

        self._collection_name = self.final_namespace

    async def initialize(self):
        async with get_data_init_lock():
            if self.db is None:
                self.db = await ClientManager.get_client()

            self._data = await get_or_create_collection(self.db, self._collection_name)

            # Create and migrate all indexes including Chinese collation for file_path
            await self.create_and_migrate_indexes_if_not_exists()

            logger.debug(
                f"[{self.workspace}] Use MongoDB as DocStatus {self._collection_name}"
            )

    async def finalize(self):
        if self.db is not None:
            await ClientManager.release_client(self.db)
            self.db = None
            self._data = None

    async def get_by_id(self, id: str) -> Union[dict[str, Any], None]:
        return await self._data.find_one({"_id": id})

    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
        cursor = self._data.find({"_id": {"$in": ids}})
        docs = await cursor.to_list(length=None)

        doc_map: dict[str, dict[str, Any]] = {}
        for doc in docs:
            if not doc:
                continue
            doc_map[str(doc.get("_id"))] = doc

        ordered_results: list[dict[str, Any] | None] = []
        for id_value in ids:
            ordered_results.append(doc_map.get(str(id_value)))
        return ordered_results

    async def filter_keys(self, data: set[str]) -> set[str]:
        cursor = self._data.find({"_id": {"$in": list(data)}}, {"_id": 1})
        existing_ids = {str(x["_id"]) async for x in cursor}
        return data - existing_ids

    async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
        logger.debug(f"[{self.workspace}] Inserting {len(data)} to {self.namespace}")
        if not data:
            return
        update_tasks: list[Any] = []
        for i, (k, v) in enumerate(data.items(), start=1):
            # Ensure chunks_list field exists and is an array
            if "chunks_list" not in v:
                v["chunks_list"] = []
            data[k]["_id"] = k
            update_tasks.append(
                self._data.update_one({"_id": k}, {"$set": v}, upsert=True)
            )
            await _cooperative_yield(i)
        await asyncio.gather(*update_tasks)

    async def get_status_counts(self) -> dict[str, int]:
        """Get counts of documents in each status"""
        pipeline = [{"$group": {"_id": "$status", "count": {"$sum": 1}}}]
        cursor = await self._data.aggregate(pipeline, allowDiskUse=True)
        result = await cursor.to_list()
        counts = {}
        for doc in result:
            counts[doc["_id"]] = doc["count"]
        return counts

    async def get_docs_by_status(
        self, status: DocStatus
    ) -> dict[str, DocProcessingStatus]:
        """Get all documents with a specific status"""
        return await self.get_docs_by_statuses([status])

    async def get_docs_by_statuses(
        self, statuses: list[DocStatus]
    ) -> dict[str, DocProcessingStatus]:
        """Get all documents matching any of the given statuses in a single query.

        Uses MongoDB's $in operator to fetch all matching statuses in one
        round-trip instead of one find() call per status.
        """
        if not statuses:
            return {}
        status_values = [s.value for s in statuses]
        cursor = self._data.find({"status": {"$in": status_values}})
        docs = await cursor.to_list(length=None)
        result = {}
        for doc in docs:
            try:
                data = self._prepare_doc_status_data(doc)
                result[doc["_id"]] = DocProcessingStatus(**data)
            except KeyError as e:
                logger.error(
                    f"[{self.workspace}] Missing required field for document {doc['_id']}: {e}"
                )
                continue
        return result

    async def get_docs_by_track_id(
        self, track_id: str
    ) -> dict[str, DocProcessingStatus]:
        """Get all documents with a specific track_id"""
        cursor = self._data.find({"track_id": track_id})
        result = await cursor.to_list()
        processed_result = {}
        for doc in result:
            try:
                data = self._prepare_doc_status_data(doc)
                processed_result[doc["_id"]] = DocProcessingStatus(**data)
            except KeyError as e:
                logger.error(
                    f"[{self.workspace}] Missing required field for document {doc['_id']}: {e}"
                )
                continue
        return processed_result

    async def index_done_callback(self) -> None:
        # Mongo handles persistence automatically
        pass

    async def is_empty(self) -> bool:
        """Check if the storage is empty for the current workspace and namespace

        Returns:
            bool: True if storage is empty, False otherwise
        """
        try:
            # Use count_documents with limit 1 for efficiency
            count = await self._data.count_documents({}, limit=1)
            return count == 0
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error checking if storage is empty: {e}")
            return True

    async def drop(self) -> dict[str, str]:
        """Drop the storage by removing all documents in the collection.

        Returns:
            dict[str, str]: Status of the operation with keys 'status' and 'message'
        """
        try:
            result = await self._data.delete_many({})
            deleted_count = result.deleted_count

            logger.info(
                f"[{self.workspace}] Dropped {deleted_count} documents from doc status {self._collection_name}"
            )
            return {
                "status": "success",
                "message": f"{deleted_count} documents dropped",
            }
        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error dropping doc status {self._collection_name}: {e}"
            )
            return {"status": "error", "message": str(e)}

    async def delete(self, ids: list[str]) -> None:
        # Convert to list if it's a set (MongoDB BSON cannot encode sets)
        if isinstance(ids, set):
            ids = list(ids)
        await self._data.delete_many({"_id": {"$in": ids}})

    async def create_and_migrate_indexes_if_not_exists(self):
        """Create indexes to optimize pagination queries and migrate file_path indexes for Chinese collation"""
        try:
            # Get indexes for the current collection only
            indexes_cursor = await self._data.list_indexes()
            existing_indexes = await indexes_cursor.to_list(length=None)
            existing_index_names = {idx.get("name", "") for idx in existing_indexes}

            # Define collation configuration for Chinese pinyin sorting
            collation_config = {"locale": "zh", "numericOrdering": True}

            # Use workspace-specific index names to avoid cross-workspace conflicts
            workspace_prefix = f"{self.workspace}_" if self.workspace != "" else ""

            # 1. Define all indexes needed with workspace-specific names
            all_indexes = [
                # Original pagination indexes
                {
                    "name": f"{workspace_prefix}status_updated_at",
                    "keys": [("status", 1), ("updated_at", -1)],
                },
                {
                    "name": f"{workspace_prefix}status_created_at",
                    "keys": [("status", 1), ("created_at", -1)],
                },
                {"name": f"{workspace_prefix}updated_at", "keys": [("updated_at", -1)]},
                {"name": f"{workspace_prefix}created_at", "keys": [("created_at", -1)]},
                {"name": f"{workspace_prefix}id", "keys": [("_id", 1)]},
                {"name": f"{workspace_prefix}track_id", "keys": [("track_id", 1)]},
                # New file_path indexes with Chinese collation and workspace-specific names
                {
                    "name": f"{workspace_prefix}file_path_zh_collation",
                    "keys": [("file_path", 1)],
                    "collation": collation_config,
                },
                {
                    "name": f"{workspace_prefix}status_file_path_zh_collation",
                    "keys": [("status", 1), ("file_path", 1)],
                    "collation": collation_config,
                },
                # Partial index on content_hash for content-based dedup lookups.
                # Mirrors the PG partial index: skip legacy/empty values so the
                # index stays small and a content_hash="" query is a guaranteed miss.
                {
                    "name": f"{workspace_prefix}content_hash",
                    "keys": [("content_hash", 1)],
                    "partialFilterExpression": {
                        "content_hash": {"$exists": True, "$type": "string", "$gt": ""}
                    },
                },
            ]

            # 2. Handle legacy index cleanup: only drop old indexes that exist in THIS collection
            legacy_index_names = [
                "file_path_zh_collation",
                "status_file_path_zh_collation",
                "status_updated_at",
                "status_created_at",
                "updated_at",
                "created_at",
                "id",
                "track_id",
                "content_hash",
            ]

            for legacy_name in legacy_index_names:
                if (
                    legacy_name in existing_index_names
                    and legacy_name
                    != f"{workspace_prefix}{legacy_name.replace(workspace_prefix, '')}"
                ):
                    try:
                        await self._data.drop_index(legacy_name)
                        logger.debug(
                            f"[{self.workspace}] Migrated: dropped legacy index '{legacy_name}' from collection {self._collection_name}"
                        )
                        existing_index_names.discard(legacy_name)
                    except PyMongoError as drop_error:
                        logger.warning(
                            f"[{self.workspace}] Failed to drop legacy index '{legacy_name}' from collection {self._collection_name}: {drop_error}"
                        )

            # 3. Create all needed indexes with workspace-specific names
            for index_info in all_indexes:
                index_name = index_info["name"]
                if index_name not in existing_index_names:
                    create_kwargs = {"name": index_name}
                    if "collation" in index_info:
                        create_kwargs["collation"] = index_info["collation"]
                    if "partialFilterExpression" in index_info:
                        create_kwargs["partialFilterExpression"] = index_info[
                            "partialFilterExpression"
                        ]

                    try:
                        await self._data.create_index(
                            index_info["keys"], **create_kwargs
                        )
                        logger.debug(
                            f"[{self.workspace}] Created index '{index_name}' for collection {self._collection_name}"
                        )
                    except PyMongoError as create_error:
                        # If creation still fails, log the error but continue with other indexes
                        logger.error(
                            f"[{self.workspace}] Failed to create index '{index_name}' for collection {self._collection_name}: {create_error}"
                        )
                else:
                    logger.debug(
                        f"[{self.workspace}] Index '{index_name}' already exists for collection {self._collection_name}"
                    )

        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error creating/migrating indexes for {self._collection_name}: {e}"
            )

    async def get_docs_paginated(
        self,
        status_filter: DocStatus | None = None,
        status_filters: list[DocStatus] | None = None,
        page: int = 1,
        page_size: int = 50,
        sort_field: str = "updated_at",
        sort_direction: str = "desc",
    ) -> tuple[list[tuple[str, DocProcessingStatus]], int]:
        """Get documents with pagination support

        Args:
            status_filter: Filter by document status, None for all statuses
            page: Page number (1-based)
            page_size: Number of documents per page (10-200)
            sort_field: Field to sort by ('created_at', 'updated_at', '_id')
            sort_direction: Sort direction ('asc' or 'desc')

        Returns:
            Tuple of (list of (doc_id, DocProcessingStatus) tuples, total_count)
        """
        status_filter_values = self.resolve_status_filter_values(
            status_filter=status_filter,
            status_filters=status_filters,
        )

        # Validate parameters
        if page < 1:
            page = 1
        if page_size < 10:
            page_size = 10
        elif page_size > 200:
            page_size = 200

        if sort_field not in ["created_at", "updated_at", "_id", "file_path"]:
            sort_field = "updated_at"

        if sort_direction.lower() not in ["asc", "desc"]:
            sort_direction = "desc"

        # Build query filter
        query_filter = {}
        if status_filter_values is not None:
            query_filter["status"] = {"$in": sorted(status_filter_values)}

        # Get total count
        total_count = await self._data.count_documents(query_filter)

        # Calculate skip value
        skip = (page - 1) * page_size

        # Build sort criteria
        sort_direction_value = 1 if sort_direction.lower() == "asc" else -1
        sort_criteria = [(sort_field, sort_direction_value)]

        # Query for paginated data with Chinese collation for file_path sorting
        if sort_field == "file_path":
            # Use Chinese collation for pinyin sorting
            cursor = (
                self._data.find(query_filter)
                .sort(sort_criteria)
                .collation({"locale": "zh", "numericOrdering": True})
                .skip(skip)
                .limit(page_size)
            )
        else:
            # Use default sorting for other fields
            cursor = (
                self._data.find(query_filter)
                .sort(sort_criteria)
                .skip(skip)
                .limit(page_size)
            )
        result = await cursor.to_list(length=page_size)

        # Convert to (doc_id, DocProcessingStatus) tuples
        documents = []
        for doc in result:
            try:
                doc_id = doc["_id"]

                data = self._prepare_doc_status_data(doc)

                doc_status = DocProcessingStatus(**data)
                documents.append((doc_id, doc_status))
            except KeyError as e:
                logger.error(
                    f"[{self.workspace}] Missing required field for document {doc['_id']}: {e}"
                )
                continue

        return documents, total_count

    async def get_all_status_counts(self) -> dict[str, int]:
        """Get counts of documents in each status for all documents

        Returns:
            Dictionary mapping status names to counts, including 'all' field
        """
        pipeline = [{"$group": {"_id": "$status", "count": {"$sum": 1}}}]
        cursor = await self._data.aggregate(pipeline, allowDiskUse=True)
        result = await cursor.to_list()

        counts = {}
        total_count = 0
        for doc in result:
            counts[doc["_id"]] = doc["count"]
            total_count += doc["count"]

        # Add 'all' field with total count
        counts["all"] = total_count

        return counts

    async def get_doc_by_file_path(self, file_path: str) -> Union[dict[str, Any], None]:
        """Get document by file path

        Args:
            file_path: The file path to search for

        Returns:
            Union[dict[str, Any], None]: Document data if found, None otherwise
            Returns the same format as get_by_id method
        """
        return await self._data.find_one({"file_path": file_path})

    async def get_doc_by_file_basename(
        self, basename: str
    ) -> Union[tuple[str, dict[str, Any]], None]:
        """Mongo-native override of basename-based document lookup.

        The caller is responsible for passing an already-canonical basename;
        stored ``file_path`` values are canonicalized by the business layer, so
        this lookup performs an exact match only and relies on the file_path
        index created by ``create_and_migrate_indexes_if_not_exists``.
        """
        if not basename:
            return None
        if basename == "unknown_source":
            return None

        try:
            doc = await self._data.find_one({"file_path": basename})
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error in get_doc_by_file_basename: {e}")
            return None
        if not doc:
            return None
        doc_id = doc.get("_id")
        if doc_id is None:
            return None
        return str(doc_id), doc

    async def get_doc_by_content_hash(
        self, content_hash: str
    ) -> Union[tuple[str, dict[str, Any]], None]:
        """Mongo-native override of content-hash document lookup.

        Uses the partial ``content_hash`` index. Empty strings are treated as a
        miss to align with the partial-index predicate; legacy rows missing the
        field cannot match a non-empty query because ``find_one`` requires an
        exact value.
        """
        if not content_hash:
            return None

        try:
            doc = await self._data.find_one({"content_hash": content_hash})
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error in get_doc_by_content_hash: {e}")
            return None
        if not doc:
            return None
        doc_id = doc.get("_id")
        if doc_id is None:
            return None
        return str(doc_id), doc


@final
@dataclass
class MongoGraphStorage(BaseGraphStorage):
    """
    A concrete implementation using MongoDB's $graphLookup to demonstrate multi-hop queries.
    """

    db: AsyncDatabase = field(default=None)
    # node collection storing node_id, node_properties
    collection: AsyncCollection = field(default=None)
    # edge collection storing source_node_id, target_node_id, and edge_properties
    edgeCollection: AsyncCollection = field(default=None)

    def __init__(self, namespace, global_config, embedding_func, workspace=None):
        super().__init__(
            namespace=namespace,
            workspace=workspace or "",
            global_config=global_config,
            embedding_func=embedding_func,
        )
        validate_workspace(self.workspace)
        # Check for MONGODB_WORKSPACE environment variable first (higher priority)
        # This allows administrators to force a specific workspace for all MongoDB storage instances
        mongodb_workspace = os.environ.get("MONGODB_WORKSPACE")
        if mongodb_workspace and mongodb_workspace.strip():
            # Use environment variable value, overriding the passed workspace parameter
            effective_workspace = mongodb_workspace.strip()
            logger.info(
                f"Using MONGODB_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
            )
        else:
            # Use the workspace parameter passed during initialization
            effective_workspace = self.workspace
            if effective_workspace:
                logger.debug(
                    f"Using passed workspace parameter: '{effective_workspace}'"
                )

        # Build final_namespace with workspace prefix for data isolation
        # Keep original namespace unchanged for type detection logic
        if effective_workspace:
            self.final_namespace = f"{effective_workspace}_{self.namespace}"
            self.workspace = effective_workspace
            logger.debug(
                f"Final namespace with workspace prefix: '{self.final_namespace}'"
            )
        else:
            # When workspace is empty, final_namespace equals original namespace
            self.final_namespace = self.namespace
            self.workspace = ""
            logger.debug(f"Final namespace (no workspace): '{self.final_namespace}'")

        self._collection_name = self.final_namespace
        self._edge_collection_name = f"{self._collection_name}_edges"
        (
            self._max_upsert_payload_bytes,
            self._max_upsert_records_per_batch,
        ) = _resolve_upsert_batch_limits()
        self._max_delete_records_per_batch = _resolve_delete_batch_limit()

    async def initialize(self):
        async with get_data_init_lock():
            if self.db is None:
                self.db = await ClientManager.get_client()

            self.collection = await get_or_create_collection(
                self.db, self._collection_name
            )
            self.edge_collection = await get_or_create_collection(
                self.db, self._edge_collection_name
            )

            # Create Atlas Search index for better search performance if possible
            await self.create_search_index_if_not_exists()

            # Fail-fast: migrate legacy edges to canonical endpoint fields and
            # build the unique index before serving (upsert relies on it). Raises
            # on failure so startup aborts rather than serving a half-migrated graph.
            await self.create_edge_indexes_and_migrate_if_not_exists()

            logger.debug(
                f"[{self.workspace}] Use MongoDB as KG {self._collection_name}"
            )

    async def finalize(self):
        if self.db is not None:
            await ClientManager.release_client(self.db)
            self.db = None
            self.collection = None
            self.edge_collection = None

    async def create_edge_indexes_and_migrate_if_not_exists(self) -> None:
        """Create the compound unique edge-endpoint index, migrating legacy edges first.

        Fail-fast one-time migration (mirrors the OpenSearch canonical-id work):

          1. dedupe legacy reciprocal duplicate docs, **merging the full relation
             payload** into the survivor (provenance unioned, keywords
             set-unioned, descriptions joined, weight summed — like
             ``_merge_edges_then_upsert``) so no relation evidence is lost;
          2. backfill the canonical ``edge_lo`` / ``edge_hi`` endpoints on every
             remaining doc;
          3. build the partial **compound** unique index on ``(edge_lo, edge_hi)``.

        The endpoints are two separate fields (not a delimiter-joined string), so
        distinct pairs never collide even if an id contains the would-be
        delimiter — no input sanitisation required.

        The index doubles as the completion flag: if it already exists we skip.
        Anything failing raises, so ``initialize``/startup aborts rather than
        serving a half-migrated collection (the upsert filter relies on every doc
        having ``edge_lo``/``edge_hi``). Runs inside ``get_data_init_lock``, so
        only the first worker of a deployment migrates; the rest skip on the index.

        Assumes no concurrent *old-version* writer adds endpoint-less docs after
        this completes (true for stop-the-world / single-deployment restarts). A
        true rolling deploy with mixed code versions writing one collection could
        leave a straggler duplicate; the remedy is to drop the
        ``edge_endpoints_unique`` index and let the next startup re-migrate.
        """
        workspace_prefix = f"{self.workspace}_" if self.workspace != "" else ""
        index_name = f"{workspace_prefix}edge_endpoints_unique"

        indexes_cursor = await self.edge_collection.list_indexes()
        existing_indexes = await indexes_cursor.to_list(length=None)
        if any(idx.get("name") == index_name for idx in existing_indexes):
            logger.info(
                f"[{self.workspace}] Edge collection {self._edge_collection_name} "
                f"already on canonical edge endpoints; skipping migration"
            )
            return

        # Best-effort total for an X/total denominator (estimated_document_count
        # is O(1) metadata); migration still works if it is unavailable.
        try:
            total = await self.edge_collection.estimated_document_count()
        except PyMongoError:
            total = None
        logger.info(
            f"[{self.workspace}] Starting canonical edge migration for "
            f"{self._edge_collection_name}"
            + (f" (~{total} edges to scan)" if total is not None else "")
        )

        removed = await self._dedupe_legacy_edges()
        backfilled = await self._backfill_edge_endpoints()
        # The unique index is the completion flag — only created on full success.
        # unique build raises if any duplicate slipped through (e.g. a concurrent
        # old-version writer), which fails startup so the next run retries.
        await self.edge_collection.create_index(
            [("edge_lo", 1), ("edge_hi", 1)],
            name=index_name,
            unique=True,
            partialFilterExpression={
                "edge_lo": {"$exists": True, "$type": "string"},
                "edge_hi": {"$exists": True, "$type": "string"},
            },
        )
        scanned = total if total is not None else "?"
        logger.info(
            f"[{self.workspace}] Canonical edge migration complete for "
            f"{self._edge_collection_name}: scanned {scanned}, deduped {removed}, "
            f"backfilled {backfilled}"
        )

    async def _dedupe_legacy_edges(self) -> int:
        """Collapse duplicate docs for the same undirected edge into one.

        Groups by the canonical (sorted) endpoint pair; for each group with more
        than one doc, keeps the newest by ``created_at`` and **merges the
        non-survivors' relation payload into it before deleting them** so no
        relation evidence is lost: ``source_ids``/``source_id``/``file_path`` and
        ``description`` are unioned over their ``GRAPH_FIELD_SEP`` components,
        ``keywords`` are comma-set-unioned, and ``weight`` is **summed** (like
        ``_merge_edges_then_upsert`` — duplicate docs carry separate accumulated
        weight).

        The merge is **idempotent across retries**: if a transient error aborts
        startup after the survivor update but before the delete, the next run
        re-processes the same group and must produce the same survivor. The union
        fields union their split components (re-merging an already-merged
        survivor is a no-op), and the weight sum counts the survivor's current
        weight once plus each other duplicate only while its source_ids are not
        yet folded into the survivor — so a retry (whose survivor already
        contains them) does not double-count. Returns the number of docs removed.
        """
        pipeline = [
            {
                "$group": {
                    "_id": {
                        "lo": {
                            "$cond": [
                                {"$lte": ["$source_node_id", "$target_node_id"]},
                                "$source_node_id",
                                "$target_node_id",
                            ]
                        },
                        "hi": {
                            "$cond": [
                                {"$lte": ["$source_node_id", "$target_node_id"]},
                                "$target_node_id",
                                "$source_node_id",
                            ]
                        },
                    },
                    "docs": {
                        "$push": {
                            "_id": "$_id",
                            "source_id": "$source_id",
                            "source_ids": "$source_ids",
                            "file_path": "$file_path",
                            "description": "$description",
                            "keywords": "$keywords",
                            "weight": "$weight",
                            "created_at": "$created_at",
                        }
                    },
                    "count": {"$sum": 1},
                }
            },
            {"$match": {"count": {"$gt": 1}}},
        ]
        removed = 0
        next_progress = _EDGE_MIGRATION_PROGRESS_INTERVAL
        cursor = await self.edge_collection.aggregate(pipeline, allowDiskUse=True)
        async for group in cursor:
            docs = group["docs"]
            survivor = max(docs, key=lambda d: d.get("created_at") or 0)
            others = [d for d in docs if d["_id"] != survivor["_id"]]
            if not others:
                continue

            # Merge the full relation payload across ALL docs (survivor included).
            # The union fields (source_ids/file_path/description/keywords) union
            # their split components, so re-merging an already-merged survivor (a
            # fail-fast retry) is a no-op.
            all_source_ids: list[str] = []
            all_file_paths: list[str] = []
            all_descriptions: list[str] = []
            all_keywords: set[str] = set()
            for d in docs:
                all_source_ids = merge_source_ids(
                    all_source_ids, _edge_source_id_list(d)
                )
                fp = d.get("file_path")
                all_file_paths = merge_source_ids(
                    all_file_paths, fp.split(GRAPH_FIELD_SEP) if fp else []
                )
                desc = d.get("description")
                all_descriptions = merge_source_ids(
                    all_descriptions, desc.split(GRAPH_FIELD_SEP) if desc else []
                )
                kw = d.get("keywords")
                if kw:
                    all_keywords.update(k.strip() for k in kw.split(",") if k.strip())

            # Weight is summed like _merge_edges_then_upsert (duplicate docs carry
            # separate accumulated evidence), but idempotently: the survivor's
            # current weight is the base (counted once) and each other duplicate
            # adds its weight ONLY if its source_ids are not already folded into
            # the survivor. On a fail-fast retry the survivor already contains the
            # others' source_ids, so they are skipped and the sum stays stable.
            # Legacy string weights are coerced; non-numeric values are skipped so
            # the migration cannot crash on a bad value.
            survivor_sids = set(_edge_source_id_list(survivor))
            weights: list[float] = []
            sw = _coerce_weight(survivor.get("weight"))
            if sw is not None:
                weights.append(sw)
            for o in others:
                o_sids = set(_edge_source_id_list(o))
                if not o_sids or o_sids <= survivor_sids:
                    continue  # no new trackable evidence -> don't (re-)add weight
                ow = _coerce_weight(o.get("weight"))
                if ow is not None:
                    weights.append(ow)

            set_fields: dict[str, Any] = {}
            if all_source_ids:
                set_fields["source_ids"] = all_source_ids
                set_fields["source_id"] = GRAPH_FIELD_SEP.join(all_source_ids)
            if all_file_paths:
                set_fields["file_path"] = GRAPH_FIELD_SEP.join(all_file_paths)
            if all_descriptions:
                set_fields["description"] = GRAPH_FIELD_SEP.join(all_descriptions)
            if all_keywords:
                set_fields["keywords"] = ",".join(sorted(all_keywords))
            if weights:
                set_fields["weight"] = sum(weights)
            if set_fields:
                await self.edge_collection.update_one(
                    {"_id": survivor["_id"]}, {"$set": set_fields}
                )
            await self.edge_collection.delete_many(
                {"_id": {"$in": [d["_id"] for d in others]}}
            )
            removed += len(others)
            if removed >= next_progress:
                logger.info(
                    f"[{self.workspace}] Canonical edge migration progress: "
                    f"deduped {removed} duplicate doc(s) so far"
                )
                next_progress += _EDGE_MIGRATION_PROGRESS_INTERVAL
        return removed

    async def _backfill_edge_endpoints(self) -> int:
        """Set the canonical ``edge_lo``/``edge_hi`` on every doc that lacks them.

        Returns the modified count. Runs after dedupe, so each canonical pair has
        one doc and the backfilled (edge_lo, edge_hi) pairs are unique.
        """
        is_sorted = {"$lte": ["$source_node_id", "$target_node_id"]}
        result = await self.edge_collection.update_many(
            {"edge_lo": {"$exists": False}},
            [
                {
                    "$set": {
                        "edge_lo": {
                            "$cond": [
                                is_sorted,
                                "$source_node_id",
                                "$target_node_id",
                            ]
                        },
                        "edge_hi": {
                            "$cond": [
                                is_sorted,
                                "$target_node_id",
                                "$source_node_id",
                            ]
                        },
                    }
                }
            ],
        )
        return result.modified_count

    # Sample entity document
    # "source_ids" is Array representation of "source_id" split by GRAPH_FIELD_SEP

    # {
    #     "_id" : "CompanyA",
    #     "entity_id" : "CompanyA",
    #     "entity_type" : "Organization",
    #     "description" : "A major technology company",
    #     "source_id" : "chunk-eeec0036b909839e8ec4fa150c939eec",
    #     "source_ids": ["chunk-eeec0036b909839e8ec4fa150c939eec"],
    #     "file_path" : "custom_kg",
    #     "created_at" : 1749904575
    # }

    # Sample relation document
    # {
    #     "_id" : ObjectId("6856ac6e7c6bad9b5470b678"), // MongoDB build-in ObjectId
    #     "description" : "CompanyA develops ProductX",
    #     "source_node_id" : "CompanyA",
    #     "target_node_id" : "ProductX",
    #     "relationship": "Develops", // To distinguish multiple same-target relations
    #     "weight" : Double("1"),
    #     "keywords" : "develop, produce",
    #     "source_id" : "chunk-eeec0036b909839e8ec4fa150c939eec",
    #     "source_ids": ["chunk-eeec0036b909839e8ec4fa150c939eec"],
    #     "file_path" : "custom_kg",
    #     "created_at" : 1749904575
    # }

    #
    # -------------------------------------------------------------------------
    # BASIC QUERIES
    # -------------------------------------------------------------------------
    #

    async def has_node(self, node_id: str) -> bool:
        """
        Check if node_id is present in the collection by looking up its doc.
        No real need for $graphLookup here, but let's keep it direct.
        """
        doc = await self.collection.find_one({"_id": node_id}, {"_id": 1})
        return doc is not None

    async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
        """
        Check if there's a direct single-hop edge between source_node_id and target_node_id.

        Matches on the canonical ``(edge_lo, edge_hi)`` pair (direction-independent)
        instead of the bidirectional ``$or``, so this point lookup is served by the
        compound unique index. Safe because the fail-fast migration in
        ``initialize`` guarantees every served doc carries the endpoints.
        """
        edge_lo, edge_hi = _canonical_edge_endpoints(source_node_id, target_node_id)
        doc = await self.edge_collection.find_one(
            {"edge_lo": edge_lo, "edge_hi": edge_hi},
            {"_id": 1},
        )
        return doc is not None

    #
    # -------------------------------------------------------------------------
    # DEGREES
    # -------------------------------------------------------------------------
    #

    async def node_degree(self, node_id: str) -> int:
        """
        Returns the total number of edges connected to node_id (both inbound and outbound).
        """
        return await self.edge_collection.count_documents(
            {"$or": [{"source_node_id": node_id}, {"target_node_id": node_id}]}
        )

    async def edge_degree(self, src_id: str, tgt_id: str) -> int:
        """Get the total degree (sum of relationships) of two nodes.

        Args:
            src_id: Label of the source node
            tgt_id: Label of the target node

        Returns:
            int: Sum of the degrees of both nodes
        """
        src_degree = await self.node_degree(src_id)
        trg_degree = await self.node_degree(tgt_id)

        return src_degree + trg_degree

    #
    # -------------------------------------------------------------------------
    # GETTERS
    # -------------------------------------------------------------------------
    #

    async def get_node(self, node_id: str) -> dict[str, str] | None:
        """
        Return the node properties, or None if missing.

        The Mongo-managed ``_id`` (which holds the entity name) is stripped so
        the returned dict carries only node properties, matching the contract
        honored by the other backends. Leaving it in lets callers that re-upsert
        a fetched node (e.g. entity rename) push ``_id`` into ``$set``, which
        MongoDB rejects as a modification of the immutable ``_id``.
        """
        doc = await self.collection.find_one({"_id": node_id})
        if doc is not None:
            doc.pop("_id", None)
        return doc

    async def get_edge(
        self, source_node_id: str, target_node_id: str
    ) -> dict[str, str] | None:
        # Canonical (edge_lo, edge_hi) point lookup served by the compound unique
        # index (see has_edge); the fail-fast migration guarantees the endpoints.
        edge_lo, edge_hi = _canonical_edge_endpoints(source_node_id, target_node_id)
        doc = await self.edge_collection.find_one(
            {"edge_lo": edge_lo, "edge_hi": edge_hi}
        )
        if doc is not None:
            # Strip the Mongo-managed ``_id`` so re-upserting a fetched edge
            # (e.g. relation rewrite during entity rename) cannot push ``_id``
            # into ``$set`` and trip the immutable-field error.
            doc.pop("_id", None)
        return doc

    async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
        """
        Retrieves all edges (relationships) for a particular node identified by its label.

        Args:
            source_node_id: Label of the node to get edges for

        Returns:
            list[tuple[str, str]]: List of (source_label, target_label) tuples representing edges
            None: If no edges found
        """
        cursor = self.edge_collection.find(
            {
                "$or": [
                    {"source_node_id": source_node_id},
                    {"target_node_id": source_node_id},
                ]
            },
            {"source_node_id": 1, "target_node_id": 1},
        )

        return [
            (e.get("source_node_id"), e.get("target_node_id")) async for e in cursor
        ]

    async def get_nodes_batch(self, node_ids: list[str]) -> dict[str, dict]:
        result = {}

        async for doc in self.collection.find({"_id": {"$in": node_ids}}):
            node_id = doc.pop("_id")
            result[node_id] = doc
        return result

    async def node_degrees_batch(self, node_ids: list[str]) -> dict[str, int]:
        # merge the outbound and inbound results with the same "_id" and sum the "degree"
        merged_results = {}

        # Outbound degrees
        outbound_pipeline = [
            {"$match": {"source_node_id": {"$in": node_ids}}},
            {"$group": {"_id": "$source_node_id", "degree": {"$sum": 1}}},
        ]

        cursor = await self.edge_collection.aggregate(
            outbound_pipeline, allowDiskUse=True
        )
        async for doc in cursor:
            merged_results[doc.get("_id")] = doc.get("degree")

        # Inbound degrees
        inbound_pipeline = [
            {"$match": {"target_node_id": {"$in": node_ids}}},
            {"$group": {"_id": "$target_node_id", "degree": {"$sum": 1}}},
        ]

        cursor = await self.edge_collection.aggregate(
            inbound_pipeline, allowDiskUse=True
        )
        async for doc in cursor:
            merged_results[doc.get("_id")] = merged_results.get(
                doc.get("_id"), 0
            ) + doc.get("degree")

        return merged_results

    async def get_nodes_edges_batch(
        self, node_ids: list[str]
    ) -> dict[str, list[tuple[str, str]]]:
        """
        Batch retrieve edges for multiple nodes.
        For each node, returns both outgoing and incoming edges to properly represent
        the undirected graph nature.

        Args:
            node_ids: List of node IDs (entity_id) for which to retrieve edges.

        Returns:
            A dictionary mapping each node ID to its list of edge tuples (source, target).
            For each node, the list includes both:
            - Outgoing edges: (queried_node, connected_node)
            - Incoming edges: (connected_node, queried_node)
        """
        result = {node_id: [] for node_id in node_ids}

        # Query outgoing edges (where node is the source)
        outgoing_cursor = self.edge_collection.find(
            {"source_node_id": {"$in": node_ids}},
            {"source_node_id": 1, "target_node_id": 1},
        )
        async for edge in outgoing_cursor:
            source = edge["source_node_id"]
            target = edge["target_node_id"]
            result[source].append((source, target))

        # Query incoming edges (where node is the target)
        incoming_cursor = self.edge_collection.find(
            {"target_node_id": {"$in": node_ids}},
            {"source_node_id": 1, "target_node_id": 1},
        )
        async for edge in incoming_cursor:
            source = edge["source_node_id"]
            target = edge["target_node_id"]
            result[target].append((source, target))

        return result

    #
    # -------------------------------------------------------------------------
    # UPSERTS
    # -------------------------------------------------------------------------
    #

    async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
        """
        Insert or update a node document.
        """
        update_doc = {"$set": {**node_data}}
        if node_data.get("source_id", ""):
            update_doc["$set"]["source_ids"] = node_data["source_id"].split(
                GRAPH_FIELD_SEP
            )

        await self.collection.update_one({"_id": node_id}, update_doc, upsert=True)

    async def upsert_edge(
        self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
    ) -> None:
        """Upsert the undirected edge between source_node_id and target_node_id.

        Matches on the canonical ``(edge_lo, edge_hi)`` endpoint pair
        (direction-independent) instead of the old bidirectional ``$or`` filter,
        so the compound unique index can reject a racing duplicate insert. If two
        writers race the first insert, the loser hits a ``DuplicateKeyError``; we
        retry once, which now matches the just-inserted doc and updates it.
        """
        # Ensure source node exists
        await self.upsert_node(source_node_id, {})

        edge_lo, edge_hi = _canonical_edge_endpoints(source_node_id, target_node_id)

        # Copy so we never mutate the caller's edge_data dict.
        set_doc: dict = {**edge_data}
        if edge_data.get("source_id", ""):
            set_doc["source_ids"] = edge_data["source_id"].split(GRAPH_FIELD_SEP)
        set_doc["source_node_id"] = source_node_id
        set_doc["target_node_id"] = target_node_id
        set_doc["edge_lo"] = edge_lo
        set_doc["edge_hi"] = edge_hi
        update_doc = {"$set": set_doc}

        for attempt in range(2):
            try:
                await self.edge_collection.update_one(
                    {"edge_lo": edge_lo, "edge_hi": edge_hi}, update_doc, upsert=True
                )
                return
            except DuplicateKeyError:
                # Another writer inserted this edge between our filter miss and
                # insert. Retry once: the doc now exists, so the upsert becomes a
                # plain update. A second failure is unexpected — let it surface.
                if attempt == 1:
                    raise

    async def upsert_nodes_batch(self, nodes: list[tuple[str, dict[str, str]]]) -> None:
        """Batch insert/update multiple nodes using a single bulk_write() call.

        Args:
            nodes: List of (node_id, node_data) tuples.
        """
        if not nodes:
            return
        ops: list[tuple[Any, int, str]] = []
        for node_id, node_data in nodes:
            update_doc: dict = {"$set": {**node_data}}
            if node_data.get("source_id", ""):
                update_doc["$set"]["source_ids"] = node_data["source_id"].split(
                    GRAPH_FIELD_SEP
                )
            ops.append(
                (
                    UpdateOne({"_id": node_id}, update_doc, upsert=True),
                    _estimate_doc_bytes(update_doc),
                    node_id,
                )
            )
        await _run_batched_bulk_write(
            self.collection,
            ops,
            max_payload_bytes=self._max_upsert_payload_bytes,
            max_records_per_batch=self._max_upsert_records_per_batch,
            ordered=True,
            log_prefix=f"[{self.workspace}] {self.namespace} nodes:",
            what="node upsert",
        )

    async def has_nodes_batch(self, node_ids: list[str]) -> set[str]:
        """Check existence of multiple nodes using a single $in query.

        Args:
            node_ids: List of node IDs to check.

        Returns:
            Set of node_ids that exist in the graph.
        """
        if not node_ids:
            return set()
        cursor = self.collection.find({"_id": {"$in": node_ids}}, {"_id": 1})
        return {doc["_id"] async for doc in cursor}

    async def upsert_edges_batch(
        self, edges: list[tuple[str, str, dict[str, str]]]
    ) -> None:
        """Batch insert/update multiple edges using a single bulk_write() call.

        Also ensures source nodes exist (matching upsert_edge() behaviour) via a
        separate bulk_write on the node collection for any source nodes that need
        to be created as empty placeholders.

        Args:
            edges: List of (source_node_id, target_node_id, edge_data) tuples.
        """
        if not edges:
            return

        # Ensure all source nodes exist (mirrors upsert_edge's upsert_node call)
        source_node_ids = list(dict.fromkeys(src for src, _tgt, _data in edges))
        node_ops: list[tuple[Any, int, str]] = [
            (
                UpdateOne({"_id": src}, {"$setOnInsert": {"_id": src}}, upsert=True),
                _estimate_doc_bytes({"_id": src}),
                src,
            )
            for src in source_node_ids
        ]
        await _run_batched_bulk_write(
            self.collection,
            node_ops,
            max_payload_bytes=self._max_upsert_payload_bytes,
            max_records_per_batch=self._max_upsert_records_per_batch,
            ordered=False,
            log_prefix=f"[{self.workspace}] {self.namespace} edges:",
            what="source-node placeholder upsert",
        )

        # Key every edge by its canonical (edge_lo, edge_hi) pair and dedupe
        # within the batch (last-write-wins). Deduping collapses reciprocal
        # directions onto one op, which both matches the compound unique index
        # and avoids an intra-batch duplicate-key error from two ops inserting
        # the same endpoint pair.
        deduped_ops: dict[tuple[str, str], tuple[Any, int, str]] = {}
        for source_node_id, target_node_id, edge_data in edges:
            update_doc: dict = {"$set": {**edge_data}}
            if edge_data.get("source_id", ""):
                update_doc["$set"]["source_ids"] = edge_data["source_id"].split(
                    GRAPH_FIELD_SEP
                )
            update_doc["$set"]["source_node_id"] = source_node_id
            update_doc["$set"]["target_node_id"] = target_node_id
            edge_lo, edge_hi = _canonical_edge_endpoints(source_node_id, target_node_id)
            update_doc["$set"]["edge_lo"] = edge_lo
            update_doc["$set"]["edge_hi"] = edge_hi
            deduped_ops[(edge_lo, edge_hi)] = (
                UpdateOne(
                    {"edge_lo": edge_lo, "edge_hi": edge_hi}, update_doc, upsert=True
                ),
                _estimate_doc_bytes(update_doc),
                f"{source_node_id}->{target_node_id}",
            )
        edge_ops = list(deduped_ops.values())

        # ordered=True (kept from the pre-canonical behaviour). Intra-batch
        # last-write-wins is already guaranteed by the endpoint-pair dedupe above
        # (one op per pair), so ordering is not load-bearing for that; we keep it
        # for continuity. If a concurrent writer (another process bypassing the keyed
        # lock) wins an insert, our upsert hits 11000 and the bulk aborts; we
        # retry the whole op list once — the racing docs now exist, so the
        # upserts update instead of inserting (idempotent). A non-11000 / write-
        # concern error re-raises rather than being masked.
        async def _run_edge_bulk() -> None:
            await _run_batched_bulk_write(
                self.edge_collection,
                edge_ops,
                max_payload_bytes=self._max_upsert_payload_bytes,
                max_records_per_batch=self._max_upsert_records_per_batch,
                ordered=True,
                log_prefix=f"[{self.workspace}] {self.namespace} edges:",
                what="edge upsert",
            )

        try:
            await _run_edge_bulk()
        except BulkWriteError as e:
            details = e.details or {}
            write_errors = details.get("writeErrors", [])
            # Retry ONLY when every failure is a duplicate-key race; a
            # writeConcern failure (durability problem, empty writeErrors) or any
            # other write error must surface, not be masked by a blind retry.
            #
            # NOTE: under ordered=True the bulk aborts at the FIRST failing op, so
            # writeErrors holds at most one entry — the all(...) check therefore
            # only inspects that first error, not the whole batch. Ops after it
            # never ran; they re-run when we retry the entire op list below. So a
            # non-11000 error hidden behind a leading 11000 is not masked — it
            # simply surfaces one retry later (the retry hits it and re-raises,
            # since by then the leading dup has resolved to a plain update).
            dup_only = (
                bool(write_errors)
                and all(we.get("code") == _DUPLICATE_KEY_CODE for we in write_errors)
                and not details.get("writeConcernErrors")
            )
            if not dup_only:
                raise
            logger.debug(
                f"[{self.workspace}] {self.namespace} edges: {len(write_errors)} "
                f"duplicate-key race(s) on edge upsert; retrying as updates"
            )
            await _run_edge_bulk()

    #
    # -------------------------------------------------------------------------
    # DELETION
    # -------------------------------------------------------------------------
    #

    async def delete_node(self, node_id: str) -> None:
        """
        1) Remove node's doc entirely.
        2) Remove inbound & outbound edges from any doc that references node_id.
        """
        # Remove all edges
        await self.edge_collection.delete_many(
            {"$or": [{"source_node_id": node_id}, {"target_node_id": node_id}]}
        )

        # Remove the node doc
        await self.collection.delete_one({"_id": node_id})

    #
    # -------------------------------------------------------------------------
    # QUERY
    # -------------------------------------------------------------------------
    #

    async def get_all_labels(self) -> list[str]:
        """
        Get all existing node _ids(entity names) in the database
        Returns:
            [id1, id2, ...]  # Alphabetically sorted id list
        """

        # Use aggregation with allowDiskUse for large datasets
        pipeline = [{"$project": {"_id": 1}}, {"$sort": {"_id": 1}}]
        cursor = await self.collection.aggregate(pipeline, allowDiskUse=True)
        labels = []
        async for doc in cursor:
            labels.append(doc["_id"])
        return labels

    def _construct_graph_node(
        self, node_id, node_data: dict[str, str]
    ) -> KnowledgeGraphNode:
        return KnowledgeGraphNode(
            id=node_id,
            labels=[node_id],
            properties={
                k: v
                for k, v in node_data.items()
                if k
                not in [
                    "_id",
                    "connected_edges",
                    "source_ids",
                    "edge_count",
                ]
            },
        )

    def _construct_graph_edge(self, edge_id: str, edge: dict[str, str]):
        return KnowledgeGraphEdge(
            id=edge_id,
            type=edge.get("relationship", ""),
            source=edge["source_node_id"],
            target=edge["target_node_id"],
            properties={
                k: v
                for k, v in edge.items()
                if k
                not in [
                    "_id",
                    "source_node_id",
                    "target_node_id",
                    "relationship",
                    "source_ids",
                    "edge_lo",
                    "edge_hi",
                ]
            },
        )

    async def _fetch_nodes_by_ids(
        self, node_ids: list[str], projection: dict[str, int] | None = None
    ) -> list[dict[str, Any]]:
        """Fetch nodes by ID while preserving the requested order."""
        if not node_ids:
            return []

        cursor = self.collection.find({"_id": {"$in": node_ids}}, projection)
        docs_by_id = {}
        async for doc in cursor:
            docs_by_id[str(doc["_id"])] = doc
        return [docs_by_id[node_id] for node_id in node_ids if node_id in docs_by_id]

    async def get_knowledge_graph_all_by_degree(
        self, max_depth: int, max_nodes: int
    ) -> KnowledgeGraph:
        """
        It's possible that the node with one or multiple relationships is retrieved,
        while its neighbor is not.  Then this node might seem like disconnected in UI.
        """

        total_node_count = await self.collection.count_documents({})
        result = KnowledgeGraph()
        seen_edges = set()

        result.is_truncated = total_node_count > max_nodes
        if result.is_truncated:
            # Get all node_ids ranked by degree if max_nodes exceeds total node count
            pipeline = [
                {"$project": {"source_node_id": 1, "_id": 0}},
                {"$group": {"_id": "$source_node_id", "degree": {"$sum": 1}}},
                {
                    "$unionWith": {
                        "coll": self._edge_collection_name,
                        "pipeline": [
                            {"$project": {"target_node_id": 1, "_id": 0}},
                            {
                                "$group": {
                                    "_id": "$target_node_id",
                                    "degree": {"$sum": 1},
                                }
                            },
                        ],
                    }
                },
                {"$group": {"_id": "$_id", "degree": {"$sum": "$degree"}}},
                {"$sort": {"degree": -1}},
                {"$limit": max_nodes},
            ]
            cursor = await self.edge_collection.aggregate(pipeline, allowDiskUse=True)

            node_ids = []
            async for doc in cursor:
                node_id = str(doc["_id"])
                node_ids.append(node_id)

            if len(node_ids) < max_nodes:
                remaining = max_nodes - len(node_ids)
                cursor = self.collection.find(
                    {"_id": {"$nin": node_ids}},
                    {"source_ids": 0},
                ).limit(remaining)
                async for doc in cursor:
                    node_ids.append(str(doc["_id"]))

            docs = await self._fetch_nodes_by_ids(node_ids, {"source_ids": 0})
            for doc in docs:
                result.nodes.append(self._construct_graph_node(doc["_id"], doc))

            # As node count reaches the limit, only need to fetch the edges that directly connect to these nodes
            edge_cursor = self.edge_collection.find(
                {
                    "$and": [
                        {"source_node_id": {"$in": node_ids}},
                        {"target_node_id": {"$in": node_ids}},
                    ]
                }
            )
        else:
            # All nodes and edges are needed
            cursor = self.collection.find({}, {"source_ids": 0})

            async for doc in cursor:
                node_id = str(doc["_id"])
                result.nodes.append(self._construct_graph_node(doc["_id"], doc))

            edge_cursor = self.edge_collection.find({})

        async for edge in edge_cursor:
            edge_id = f"{edge['source_node_id']}-{edge['target_node_id']}"
            if edge_id not in seen_edges:
                seen_edges.add(edge_id)
                result.edges.append(self._construct_graph_edge(edge_id, edge))

        return result

    async def _bidirectional_bfs_nodes(
        self,
        node_labels: list[str],
        seen_nodes: set[str],
        result: KnowledgeGraph,
        depth: int,
        max_depth: int,
        max_nodes: int,
    ) -> KnowledgeGraph:
        if depth > max_depth or len(result.nodes) > max_nodes:
            return result

        cursor = self.collection.find({"_id": {"$in": node_labels}})

        async for node in cursor:
            node_id = node["_id"]
            if node_id not in seen_nodes:
                seen_nodes.add(node_id)
                result.nodes.append(self._construct_graph_node(node_id, node))
                if len(result.nodes) > max_nodes:
                    return result

        # Collect neighbors
        # Get both inbound and outbound one hop nodes
        cursor = self.edge_collection.find(
            {
                "$or": [
                    {"source_node_id": {"$in": node_labels}},
                    {"target_node_id": {"$in": node_labels}},
                ]
            }
        )

        neighbor_nodes = []
        async for edge in cursor:
            if edge["source_node_id"] not in seen_nodes:
                neighbor_nodes.append(edge["source_node_id"])
            if edge["target_node_id"] not in seen_nodes:
                neighbor_nodes.append(edge["target_node_id"])

        if neighbor_nodes:
            result = await self._bidirectional_bfs_nodes(
                neighbor_nodes, seen_nodes, result, depth + 1, max_depth, max_nodes
            )

        return result

    async def get_knowledge_subgraph_bidirectional_bfs(
        self,
        node_label: str,
        depth: int,
        max_depth: int,
        max_nodes: int,
    ) -> KnowledgeGraph:
        seen_nodes = set()
        seen_edges = set()
        result = KnowledgeGraph()

        result = await self._bidirectional_bfs_nodes(
            [node_label], seen_nodes, result, depth, max_depth, max_nodes
        )

        # Get all edges from seen_nodes
        all_node_ids = list(seen_nodes)
        cursor = self.edge_collection.find(
            {
                "$and": [
                    {"source_node_id": {"$in": all_node_ids}},
                    {"target_node_id": {"$in": all_node_ids}},
                ]
            }
        )

        async for edge in cursor:
            edge_id = f"{edge['source_node_id']}-{edge['target_node_id']}"
            if edge_id not in seen_edges:
                result.edges.append(self._construct_graph_edge(edge_id, edge))
                seen_edges.add(edge_id)

        return result

    async def get_knowledge_subgraph_in_out_bound_bfs(
        self, node_label: str, max_depth: int, max_nodes: int
    ) -> KnowledgeGraph:
        seen_nodes = set()
        seen_edges = set()
        result = KnowledgeGraph()
        project_doc = {
            "source_ids": 0,
            "created_at": 0,
            "entity_type": 0,
            "file_path": 0,
        }

        # Verify if starting node exists
        start_node = await self.collection.find_one({"_id": node_label})
        if not start_node:
            logger.warning(
                f"[{self.workspace}] Starting node with label {node_label} does not exist!"
            )
            return result

        seen_nodes.add(node_label)
        result.nodes.append(self._construct_graph_node(node_label, start_node))

        if max_depth == 0:
            return result

        # In MongoDB, depth = 0 means one-hop
        max_depth = max_depth - 1

        pipeline = [
            {"$match": {"_id": node_label}},
            {"$project": project_doc},
            {
                "$graphLookup": {
                    "from": self._edge_collection_name,
                    "startWith": "$_id",
                    "connectFromField": "target_node_id",
                    "connectToField": "source_node_id",
                    "maxDepth": max_depth,
                    "depthField": "depth",
                    "as": "connected_edges",
                },
            },
            {
                "$unionWith": {
                    "coll": self._collection_name,
                    "pipeline": [
                        {"$match": {"_id": node_label}},
                        {"$project": project_doc},
                        {
                            "$graphLookup": {
                                "from": self._edge_collection_name,
                                "startWith": "$_id",
                                "connectFromField": "source_node_id",
                                "connectToField": "target_node_id",
                                "maxDepth": max_depth,
                                "depthField": "depth",
                                "as": "connected_edges",
                            }
                        },
                    ],
                }
            },
        ]

        cursor = await self.collection.aggregate(pipeline, allowDiskUse=True)
        node_edges = []

        # Two records for node_label are returned capturing outbound and inbound connected_edges
        async for doc in cursor:
            if doc.get("connected_edges", []):
                node_edges.extend(doc.get("connected_edges"))

        # Sort the connected edges by depth ascending and weight descending
        # And stores the source_node_id and target_node_id in sequence to retrieve the neighbouring nodes
        node_edges = sorted(
            node_edges,
            key=lambda x: (x["depth"], -x["weight"]),
        )

        # As order matters, we need to use another list to store the node_id
        # And only take the first max_nodes ones
        node_ids = []
        for edge in node_edges:
            if len(node_ids) < max_nodes and edge["source_node_id"] not in seen_nodes:
                node_ids.append(edge["source_node_id"])
                seen_nodes.add(edge["source_node_id"])

            if len(node_ids) < max_nodes and edge["target_node_id"] not in seen_nodes:
                node_ids.append(edge["target_node_id"])
                seen_nodes.add(edge["target_node_id"])

        # Filter out all the node whose id is same as node_label so that we do not check existence next step
        cursor = self.collection.find({"_id": {"$in": node_ids}})

        async for doc in cursor:
            result.nodes.append(self._construct_graph_node(str(doc["_id"]), doc))

        for edge in node_edges:
            if (
                edge["source_node_id"] not in seen_nodes
                or edge["target_node_id"] not in seen_nodes
            ):
                continue

            edge_id = f"{edge['source_node_id']}-{edge['target_node_id']}"
            if edge_id not in seen_edges:
                result.edges.append(self._construct_graph_edge(edge_id, edge))
                seen_edges.add(edge_id)

        return result

    async def get_knowledge_graph(
        self,
        node_label: str,
        max_depth: int = 3,
        max_nodes: int = None,
    ) -> KnowledgeGraph:
        """
        Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.

        Args:
            node_label: Label of the starting node, * means all nodes
            max_depth: Maximum depth of the subgraph, Defaults to 3
            max_nodes: Maximum nodes to return, Defaults to global_config max_graph_nodes

        Returns:
            KnowledgeGraph object containing nodes and edges, with an is_truncated flag
            indicating whether the graph was truncated due to max_nodes limit

        If a graph is like this and starting from B:
        A → B ← C ← F, B -> E, C → D

        Outbound BFS:
        B → E

        Inbound BFS:
        A → B
        C → B
        F → C

        Bidirectional BFS:
        A → B
        B → E
        F → C
        C → B
        C → D
        """
        # Use global_config max_graph_nodes as default if max_nodes is None
        if max_nodes is None:
            max_nodes = self.global_config.get("max_graph_nodes", 1000)
        else:
            # Limit max_nodes to not exceed global_config max_graph_nodes
            max_nodes = min(max_nodes, self.global_config.get("max_graph_nodes", 1000))

        result = KnowledgeGraph()
        start = time.perf_counter()

        try:
            # Optimize pipeline to avoid memory issues with large datasets
            if node_label == "*":
                result = await self.get_knowledge_graph_all_by_degree(
                    max_depth, max_nodes
                )
            elif GRAPH_BFS_MODE == "in_out_bound":
                result = await self.get_knowledge_subgraph_in_out_bound_bfs(
                    node_label, max_depth, max_nodes
                )
            else:
                result = await self.get_knowledge_subgraph_bidirectional_bfs(
                    node_label, 0, max_depth, max_nodes
                )

            duration = time.perf_counter() - start

            logger.info(
                f"[{self.workspace}] Subgraph query successful in {duration:.4f} seconds | Node count: {len(result.nodes)} | Edge count: {len(result.edges)} | Truncated: {result.is_truncated}"
            )

        except PyMongoError as e:
            # Handle memory limit errors specifically
            if "memory limit" in str(e).lower() or "sort exceeded" in str(e).lower():
                logger.warning(
                    f"[{self.workspace}] MongoDB memory limit exceeded, falling back to simple query: {str(e)}"
                )
                # Fallback to a simple query without complex aggregation
                try:
                    simple_cursor = self.collection.find({}).limit(max_nodes)
                    async for doc in simple_cursor:
                        result.nodes.append(
                            self._construct_graph_node(str(doc["_id"]), doc)
                        )
                    result.is_truncated = True
                    logger.info(
                        f"[{self.workspace}] Fallback query completed | Node count: {len(result.nodes)}"
                    )
                except PyMongoError as fallback_error:
                    logger.error(
                        f"[{self.workspace}] Fallback query also failed: {str(fallback_error)}"
                    )
            else:
                logger.error(f"[{self.workspace}] MongoDB query failed: {str(e)}")

        return result

    async def index_done_callback(self) -> None:
        # Mongo handles persistence automatically
        pass

    async def remove_nodes(self, nodes: list[str]) -> None:
        """Delete multiple nodes

        Args:
            nodes: List of node IDs to be deleted
        """
        logger.info(f"[{self.workspace}] Deleting {len(nodes)} nodes")
        if not nodes:
            return

        # 1. Remove all edges referencing these nodes
        await self.edge_collection.delete_many(
            {
                "$or": [
                    {"source_node_id": {"$in": nodes}},
                    {"target_node_id": {"$in": nodes}},
                ]
            }
        )

        # 2. Delete the node documents
        await self.collection.delete_many({"_id": {"$in": nodes}})

        logger.debug(f"[{self.workspace}] Successfully deleted nodes: {nodes}")

    async def remove_edges(self, edges: list[tuple[str, str]]) -> None:
        """Delete multiple edges

        Args:
            edges: List of edges to be deleted, each edge is a (source, target) tuple
        """
        logger.info(f"[{self.workspace}] Deleting {len(edges)} edges")
        if not edges:
            return

        # Match each edge by its canonical (edge_lo, edge_hi) pair: one clause per
        # edge (vs. the old two-clause bidirectional pair) served by the compound
        # unique index, with reciprocal/duplicate inputs collapsed. Safe because
        # the fail-fast migration guarantees every served doc carries the endpoints.
        seen: set[tuple[str, str]] = set()
        all_edge_pairs = []
        for source_id, target_id in edges:
            endpoints = _canonical_edge_endpoints(source_id, target_id)
            if endpoints in seen:
                continue
            seen.add(endpoints)
            all_edge_pairs.append({"edge_lo": endpoints[0], "edge_hi": endpoints[1]})

        # Chunk the $or by record count so a large delete stays under the bulk
        # message / 16MB query limit; endpoints are bounded id strings, so a count
        # cap is enough (no byte budget needed). A non-positive cap disables it.
        chunk = (
            self._max_delete_records_per_batch
            if self._max_delete_records_per_batch > 0
            else len(all_edge_pairs)
        )
        for i in range(0, len(all_edge_pairs), chunk):
            await self.edge_collection.delete_many(
                {"$or": all_edge_pairs[i : i + chunk]}
            )

        logger.debug(f"[{self.workspace}] Successfully deleted edges: {edges}")

    async def get_all_nodes(self) -> list[dict]:
        """Get all nodes in the graph.

        Returns:
            A list of all nodes, where each node is a dictionary of its properties
        """
        cursor = self.collection.find({})
        nodes = []
        async for node in cursor:
            node_dict = dict(node)
            # Add node id (entity_id) to the dictionary for easier access
            node_dict["id"] = node_dict.get("_id")
            nodes.append(node_dict)
        return nodes

    async def get_all_edges(self) -> list[dict]:
        """Get all edges in the graph.

        Returns:
            A list of all edges, where each edge is a dictionary of its properties
        """
        cursor = self.edge_collection.find({})
        edges = []
        async for edge in cursor:
            edge_dict = dict(edge)
            edge_dict["source"] = edge_dict.get("source_node_id")
            edge_dict["target"] = edge_dict.get("target_node_id")
            edges.append(edge_dict)
        return edges

    async def get_popular_labels(self, limit: int = 300) -> list[str]:
        """Get popular labels(entity names) by node degree (most connected entities)

        Args:
            limit: Maximum number of labels to return

        Returns:
            List of labels(entity names) sorted by degree (highest first)
        """
        try:
            # Use aggregation pipeline to count edges per node and sort by degree
            pipeline = [
                # Count outbound edges
                {"$group": {"_id": "$source_node_id", "out_degree": {"$sum": 1}}},
                # Union with inbound edges count
                {
                    "$unionWith": {
                        "coll": self._edge_collection_name,
                        "pipeline": [
                            {
                                "$group": {
                                    "_id": "$target_node_id",
                                    "in_degree": {"$sum": 1},
                                }
                            }
                        ],
                    }
                },
                # Group by node_id and sum degrees
                {
                    "$group": {
                        "_id": "$_id",
                        "total_degree": {
                            "$sum": {
                                "$add": [
                                    {"$ifNull": ["$out_degree", 0]},
                                    {"$ifNull": ["$in_degree", 0]},
                                ]
                            }
                        },
                    }
                },
                # Sort by degree descending, then by label ascending
                {"$sort": {"total_degree": -1, "_id": 1}},
                # Limit results
                {"$limit": limit},
                # Project only the label
                {"$project": {"_id": 1}},
            ]

            cursor = await self.edge_collection.aggregate(pipeline, allowDiskUse=True)
            labels = []
            async for doc in cursor:
                if doc.get("_id"):
                    labels.append(doc["_id"])

            logger.debug(
                f"[{self.workspace}] Retrieved {len(labels)} popular labels (limit: {limit})"
            )
            return labels
        except Exception as e:
            logger.error(f"[{self.workspace}] Error getting popular labels: {str(e)}")
            return []

    async def _try_atlas_text_search(self, query_strip: str, limit: int) -> list[str]:
        """Try Atlas Search using simple text search."""
        try:
            pipeline = [
                {
                    "$search": {
                        "index": "entity_id_search_idx",
                        "text": {"query": query_strip, "path": "_id"},
                    }
                },
                {"$project": {"_id": 1, "score": {"$meta": "searchScore"}}},
                {"$limit": limit},
            ]
            cursor = await self.collection.aggregate(pipeline)
            labels = [doc["_id"] async for doc in cursor if doc.get("_id")]
            if labels:
                logger.debug(
                    f"[{self.workspace}] Atlas text search returned {len(labels)} results"
                )
                return labels
            return []
        except PyMongoError as e:
            logger.debug(f"[{self.workspace}] Atlas text search failed: {e}")
            return []

    async def _try_atlas_autocomplete_search(
        self, query_strip: str, limit: int
    ) -> list[str]:
        """Try Atlas Search using autocomplete for prefix matching."""
        try:
            pipeline = [
                {
                    "$search": {
                        "index": "entity_id_search_idx",
                        "autocomplete": {
                            "query": query_strip,
                            "path": "_id",
                            "fuzzy": {"maxEdits": 1, "prefixLength": 1},
                        },
                    }
                },
                {"$project": {"_id": 1, "score": {"$meta": "searchScore"}}},
                {"$limit": limit},
            ]
            cursor = await self.collection.aggregate(pipeline)
            labels = [doc["_id"] async for doc in cursor if doc.get("_id")]
            if labels:
                logger.debug(
                    f"[{self.workspace}] Atlas autocomplete search returned {len(labels)} results"
                )
                return labels
            return []
        except PyMongoError as e:
            logger.debug(f"[{self.workspace}] Atlas autocomplete search failed: {e}")
            return []

    async def _try_atlas_compound_search(
        self, query_strip: str, limit: int
    ) -> list[str]:
        """Try Atlas Search using compound query for comprehensive matching."""
        try:
            pipeline = [
                {
                    "$search": {
                        "index": "entity_id_search_idx",
                        "compound": {
                            "should": [
                                {
                                    "text": {
                                        "query": query_strip,
                                        "path": "_id",
                                        "score": {"boost": {"value": 10}},
                                    }
                                },
                                {
                                    "autocomplete": {
                                        "query": query_strip,
                                        "path": "_id",
                                        "score": {"boost": {"value": 5}},
                                        "fuzzy": {"maxEdits": 1, "prefixLength": 1},
                                    }
                                },
                                {
                                    "wildcard": {
                                        "query": f"*{query_strip}*",
                                        "path": "_id",
                                        "score": {"boost": {"value": 2}},
                                    }
                                },
                            ],
                            "minimumShouldMatch": 1,
                        },
                    }
                },
                {"$project": {"_id": 1, "score": {"$meta": "searchScore"}}},
                {"$sort": {"score": {"$meta": "searchScore"}}},
                {"$limit": limit},
            ]
            cursor = await self.collection.aggregate(pipeline)
            labels = [doc["_id"] async for doc in cursor if doc.get("_id")]
            if labels:
                logger.debug(
                    f"[{self.workspace}] Atlas compound search returned {len(labels)} results"
                )
                return labels
            return []
        except PyMongoError as e:
            logger.debug(f"[{self.workspace}] Atlas compound search failed: {e}")
            return []

    async def _fallback_regex_search(self, query_strip: str, limit: int) -> list[str]:
        """Fallback to regex-based search when Atlas Search fails."""
        try:
            logger.debug(
                f"[{self.workspace}] Using regex fallback search for: '{query_strip}'"
            )

            escaped_query = re.escape(query_strip)
            regex_condition = {"_id": {"$regex": escaped_query, "$options": "i"}}
            cursor = self.collection.find(regex_condition, {"_id": 1}).limit(limit * 2)
            docs = await cursor.to_list(length=limit * 2)

            # Extract labels
            labels = []
            for doc in docs:
                doc_id = doc.get("_id")
                if doc_id:
                    labels.append(doc_id)

            # Sort results to prioritize exact matches and starts-with matches
            def sort_key(label):
                label_lower = label.lower()
                query_lower_strip = query_strip.lower()

                if label_lower == query_lower_strip:
                    return (0, label_lower)  # Exact match - highest priority
                elif label_lower.startswith(query_lower_strip):
                    return (1, label_lower)  # Starts with - medium priority
                else:
                    return (2, label_lower)  # Contains - lowest priority

            labels.sort(key=sort_key)
            labels = labels[:limit]  # Apply final limit after sorting

            logger.debug(
                f"[{self.workspace}] Regex fallback search returned {len(labels)} results (limit: {limit})"
            )
            return labels

        except Exception as e:
            logger.error(f"[{self.workspace}] Regex fallback search failed: {e}")
            import traceback

            logger.error(f"[{self.workspace}] Traceback: {traceback.format_exc()}")
            return []

    async def search_labels(self, query: str, limit: int = 50) -> list[str]:
        """
        Search labels(entity names) with progressive fallback strategy:
        1. Atlas text search (simple and fast)
        2. Atlas autocomplete search (prefix matching with fuzzy)
        3. Atlas compound search (comprehensive matching)
        4. Regex fallback (when Atlas Search is unavailable)
        """
        query_strip = query.strip()
        if not query_strip:
            return []

        # First check if we have any nodes at all
        try:
            node_count = await self.collection.count_documents({})
            if node_count == 0:
                logger.debug(
                    f"[{self.workspace}] No nodes found in collection {self._collection_name}"
                )
                return []
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error counting nodes: {e}")
            return []

        # Progressive search strategy
        search_methods = [
            ("text", self._try_atlas_text_search),
            ("autocomplete", self._try_atlas_autocomplete_search),
            ("compound", self._try_atlas_compound_search),
        ]

        # Try Atlas Search methods in order
        for method_name, search_method in search_methods:
            try:
                labels = await search_method(query_strip, limit)
                if labels:
                    logger.debug(
                        f"[{self.workspace}] Search successful using {method_name} method: {len(labels)} results"
                    )
                    return labels
                else:
                    logger.debug(
                        f"[{self.workspace}] {method_name} search returned no results, trying next method"
                    )
            except Exception as e:
                logger.debug(
                    f"[{self.workspace}] {method_name} search failed: {e}, trying next method"
                )
                continue

        # If all Atlas Search methods fail, use regex fallback
        logger.info(
            f"[{self.workspace}] All Atlas Search methods failed, using regex fallback search for: '{query_strip}'"
        )
        return await self._fallback_regex_search(query_strip, limit)

    async def _check_if_index_needs_rebuild(
        self, indexes: list, index_name: str
    ) -> bool:
        """Check if the existing index needs to be rebuilt due to configuration issues."""
        for index in indexes:
            if index["name"] == index_name:
                # Check if the index has the old problematic configuration
                definition = index.get("latestDefinition", {})
                mappings = definition.get("mappings", {})
                fields = mappings.get("fields", {})
                id_field = fields.get("_id", {})

                # If it's the old single-type autocomplete configuration, rebuild
                if (
                    isinstance(id_field, dict)
                    and id_field.get("type") == "autocomplete"
                ):
                    logger.info(
                        f"[{self.workspace}] Found old index configuration for '{index_name}', will rebuild"
                    )
                    return True

                # If it's not a list (multi-type configuration), rebuild
                if not isinstance(id_field, list):
                    logger.info(
                        f"[{self.workspace}] Index '{index_name}' needs upgrade to multi-type configuration"
                    )
                    return True

                logger.info(
                    f"[{self.workspace}] Index '{index_name}' has correct configuration"
                )
                return False
        return True  # Index doesn't exist, needs creation

    async def _safely_drop_old_index(self, index_name: str):
        """Safely drop the old search index."""
        try:
            await self.collection.drop_search_index(index_name)
            logger.info(
                f"[{self.workspace}] Successfully dropped old search index '{index_name}'"
            )
        except PyMongoError as e:
            logger.warning(
                f"[{self.workspace}] Could not drop old index '{index_name}': {e}"
            )

    async def _create_improved_search_index(self, index_name: str):
        """Create an improved search index with multiple field types."""
        search_index_model = SearchIndexModel(
            definition={
                "mappings": {
                    "dynamic": False,
                    "fields": {
                        "_id": [
                            {
                                "type": "string",
                            },
                            {
                                "type": "token",
                            },
                            {
                                "type": "autocomplete",
                                "maxGrams": 15,
                                "minGrams": 2,
                            },
                        ]
                    },
                },
                "analyzer": "lucene.standard",  # Index-level analyzer for text processing
            },
            name=index_name,
            type="search",
        )

        await self.collection.create_search_index(search_index_model)
        logger.info(
            f"[{self.workspace}] Created improved Atlas Search index '{index_name}' for collection {self._collection_name}. "
        )
        logger.info(
            f"[{self.workspace}] Index will be built asynchronously, using regex fallback until ready."
        )

    async def create_search_index_if_not_exists(self):
        """Creates an improved Atlas Search index for entity search, rebuilding if necessary."""
        index_name = "entity_id_search_idx"

        try:
            # Check if we're using MongoDB Atlas (has search index capabilities)
            indexes_cursor = await self.collection.list_search_indexes()
            indexes = await indexes_cursor.to_list(length=None)

            # Check if we need to rebuild the index
            needs_rebuild = await self._check_if_index_needs_rebuild(
                indexes, index_name
            )

            if needs_rebuild:
                # Check if index exists and drop it
                index_exists = any(idx["name"] == index_name for idx in indexes)
                if index_exists:
                    await self._safely_drop_old_index(index_name)

                # Create the improved search index (async, no waiting)
                await self._create_improved_search_index(index_name)
            else:
                logger.info(
                    f"[{self.workspace}] Atlas Search index '{index_name}' already exists with correct configuration"
                )

        except PyMongoError as e:
            # This is expected if not using MongoDB Atlas or if search indexes are not supported
            logger.info(
                f"[{self.workspace}] Could not create Atlas Search index for {self._collection_name}: {e}. "
                "This is normal if not using MongoDB Atlas - search will use regex fallback."
            )
        except Exception as e:
            logger.warning(
                f"[{self.workspace}] Unexpected error creating Atlas Search index for {self._collection_name}: {e}"
            )

    async def drop(self) -> dict[str, str]:
        """Drop the storage by removing all documents in the collection.

        Returns:
            dict[str, str]: Status of the operation with keys 'status' and 'message'
        """
        try:
            result = await self.collection.delete_many({})
            deleted_count = result.deleted_count

            logger.info(
                f"[{self.workspace}] Dropped {deleted_count} documents from graph {self._collection_name}"
            )

            result = await self.edge_collection.delete_many({})
            edge_count = result.deleted_count
            logger.info(
                f"[{self.workspace}] Dropped {edge_count} edges from graph {self._edge_collection_name}"
            )

            return {
                "status": "success",
                "message": f"{deleted_count} documents and {edge_count} edges dropped",
            }
        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error dropping graph {self._collection_name}: {e}"
            )
            return {"status": "error", "message": str(e)}


@dataclass
class _PendingVectorDoc:
    """Buffered vector upsert waiting for embedding and/or bulk flush."""

    source: dict[str, Any]
    content: str
    vector: list[float] | None = None


@final
@dataclass
class MongoVectorDBStorage(BaseVectorStorage):
    db: AsyncDatabase | None = field(default=None)
    _data: AsyncCollection | None = field(default=None)
    _index_name: str = field(default="", init=False)

    def __init__(
        self, namespace, global_config, embedding_func, workspace=None, meta_fields=None
    ):
        super().__init__(
            namespace=namespace,
            workspace=workspace or "",
            global_config=global_config,
            embedding_func=embedding_func,
            meta_fields=meta_fields or set(),
        )
        self.__post_init__()

    def __post_init__(self):
        validate_workspace(self.workspace)
        self._validate_embedding_func()

        # Check for MONGODB_WORKSPACE environment variable first (higher priority)
        # This allows administrators to force a specific workspace for all MongoDB storage instances
        mongodb_workspace = os.environ.get("MONGODB_WORKSPACE")
        if mongodb_workspace and mongodb_workspace.strip():
            # Use environment variable value, overriding the passed workspace parameter
            effective_workspace = mongodb_workspace.strip()
            logger.info(
                f"Using MONGODB_WORKSPACE environment variable: '{effective_workspace}' (overriding '{self.workspace}/{self.namespace}')"
            )
        else:
            # Use the workspace parameter passed during initialization
            effective_workspace = self.workspace
            if effective_workspace:
                logger.debug(
                    f"Using passed workspace parameter: '{effective_workspace}'"
                )

        # Build final_namespace with workspace prefix for data isolation
        # Keep original namespace unchanged for type detection logic
        if effective_workspace:
            self.final_namespace = f"{effective_workspace}_{self.namespace}"
            self.workspace = effective_workspace
            logger.debug(
                f"Final namespace with workspace prefix: '{self.final_namespace}'"
            )
        else:
            # When workspace is empty, final_namespace equals original namespace
            self.final_namespace = self.namespace
            self.workspace = ""
            logger.debug(f"Final namespace (no workspace): '{self.final_namespace}'")

        # Set index name based on workspace for backward compatibility
        if effective_workspace:
            # Use collection-specific index name for workspaced collections to avoid conflicts
            self._index_name = f"vector_knn_index_{self.final_namespace}"
        else:
            # Keep original index name for backward compatibility with existing deployments
            self._index_name = "vector_knn_index"

        kwargs = self.global_config.get("vector_db_storage_cls_kwargs", {})
        cosine_threshold = kwargs.get("cosine_better_than_threshold")
        if cosine_threshold is None:
            raise ValueError(
                "cosine_better_than_threshold must be specified in vector_db_storage_cls_kwargs"
            )
        self.cosine_better_than_threshold = cosine_threshold
        self._collection_name = self.final_namespace
        self._max_batch_size = self.global_config["embedding_batch_num"]

        # Flush-time batching limits (see module-level DEFAULT_MONGO_* constants).
        # A non-positive value disables that splitting dimension. The upsert and
        # delete caps are shared across KV/graph/VDB via the _resolve_* helpers so
        # every path stays under the same bulk message / 16MB query limit.
        (
            self._max_upsert_payload_bytes,
            self._max_upsert_records_per_batch,
        ) = _resolve_upsert_batch_limits()
        self._max_delete_records_per_batch = _resolve_delete_batch_limit()

        # Deferred-embedding buffers and the per-namespace flush lock.
        # Constructed in initialize() once shared-storage primitives are
        # available; keyed on final_namespace so two instances pointing at
        # the same MongoDB collection (e.g. with the MONGODB_WORKSPACE env
        # override) share a single writer lock.
        self._pending_vector_docs: dict[str, _PendingVectorDoc] = {}
        self._pending_vector_deletes: set[str] = set()
        self._flush_lock = None

    async def initialize(self):
        async with get_data_init_lock():
            if self.db is None:
                self.db = await ClientManager.get_client()

            self._data = await get_or_create_collection(self.db, self._collection_name)

            # Ensure vector index exists
            await self.create_vector_index_if_not_exists()

            logger.debug(
                f"[{self.workspace}] Use MongoDB as VDB {self._collection_name}"
            )

        if self._flush_lock is None:
            self._flush_lock = get_namespace_lock(
                namespace=self.final_namespace, workspace=""
            )

    async def finalize(self):
        """Flush pending vector ops, release the Mongo client, surface unflushed data."""
        flush_error: Exception | None = None
        try:
            await self._flush_pending_vector_ops()
        except Exception as e:
            flush_error = e

        if self.db is not None:
            await ClientManager.release_client(self.db)
            self.db = None
            self._data = None

        pending_docs = len(self._pending_vector_docs)
        pending_deletes = len(self._pending_vector_deletes)

        if flush_error is not None:
            raise RuntimeError(
                f"[{self.workspace}] MongoVectorDBStorage.finalize() flush raised; "
                f"{pending_docs} pending upserts and {pending_deletes} pending "
                f"deletes were left buffered (client released, data lost)"
            ) from flush_error
        if pending_docs or pending_deletes:
            raise RuntimeError(
                f"[{self.workspace}] MongoVectorDBStorage.finalize() left "
                f"{pending_docs} pending upserts and {pending_deletes} pending "
                f"deletes buffered after final flush attempt (these writes have been lost)"
            )

    async def _wait_for_search_index_absent(
        self, index_name: str, *, timeout: float = 120.0, interval: float = 2.0
    ) -> None:
        """Poll until a dropped search index disappears.

        ``create_search_index`` rejects a name that still exists while the
        prior drop is in the DELETING state, so a recreate must wait for the
        old index to clear first. Best-effort: on timeout it logs and returns
        so the subsequent create surfaces any genuine conflict itself rather
        than blocking initialize() indefinitely.
        """
        deadline = time.monotonic() + timeout
        while True:
            cursor = await self._data.list_search_indexes()
            names = {idx["name"] for idx in await cursor.to_list(length=None)}
            if index_name not in names:
                return
            if time.monotonic() >= deadline:
                logger.warning(
                    f"[{self.workspace}] dropped search index {index_name} still "
                    f"present after {timeout:.0f}s; proceeding to recreate"
                )
                return
            await asyncio.sleep(interval)

    async def create_vector_index_if_not_exists(self):
        """Create the Atlas Vector Search index, repairing a FAILED one.

        Atlas/mongot leaves a vector index in the terminal ``FAILED`` state
        after a build error and never retries it on its own; when that index
        is also non-queryable every subsequent ``$vectorSearch`` raises
        ``cannot query vector index ... while in state FAILED``. Matching the
        index only by name would treat that dead index as healthy and wedge
        all queries permanently, so a non-queryable, same-dimension FAILED
        index is dropped and rebuilt here.

        Two guards run *before* the rebuild: (1) a FAILED index that is still
        ``queryable`` (a background rebuild/update failed but the previously
        built index keeps serving) is left in place to avoid taking a
        still-serving index offline; (2) a FAILED index built under a
        different embedding model raises rather than being auto-rebuilt
        against incompatible stored vectors. Transitional states
        (``PENDING``/``BUILDING``) are left alone -- they become queryable
        without intervention.
        """
        try:
            indexes_cursor = await self._data.list_search_indexes()
            indexes = await indexes_cursor.to_list(length=None)
            for index in indexes:
                if index["name"] != self._index_name:
                    continue

                # Read the stored vector dimension first so the mismatch
                # guard below runs even for a FAILED index. A FAILED index
                # built under a *different* embedding model must NOT be
                # silently auto-rebuilt: recreating with the new dimension
                # against incompatible stored vectors would just FAIL again
                # and hide the required data-directory reset from the
                # operator. Only a same-dimension FAILED index is self-healed.
                existing_dim = None
                definition = index.get("latestDefinition", {})
                fields = definition.get("fields", [])
                for field in fields:
                    if field.get("type") == "vector" and field.get("path") == "vector":
                        existing_dim = field.get("numDimensions")
                        break

                expected_dim = self.embedding_func.embedding_dim

                if existing_dim is not None and existing_dim != expected_dim:
                    error_msg = (
                        f"Vector dimension mismatch! Index '{self._index_name}' has "
                        f"dimension {existing_dim}, but current embedding model expects "
                        f"dimension {expected_dim}. Please drop the existing index or "
                        f"use an embedding model with matching dimensions."
                    )
                    logger.error(f"[{self.workspace}] {error_msg}")
                    raise ValueError(error_msg)

                # Self-heal a FAILED index, but ONLY when it is actually
                # non-queryable. Atlas can report status="FAILED" while
                # queryable=true -- e.g. a background rebuild/update failed
                # yet the previously-built index keeps serving queries (see
                # the listSearchIndexes status docs). Dropping such an index
                # here would take a still-serving index offline and cause
                # avoidable query downtime while we wait for deletion and
                # rebuild. Reached only once the dimension guard above
                # confirmed the stored dimension matches.
                if index.get("status") == "FAILED":
                    if index.get("queryable", True):
                        logger.warning(
                            f"[{self.workspace}] vector index {self._index_name} reports "
                            f"FAILED status but is still queryable; leaving the active "
                            f"index in place. A background rebuild/update likely failed -- "
                            f"inspect $listSearchIndexes statusDetail and drop/rebuild "
                            f"manually if queries degrade."
                        )
                        return

                    # Non-queryable FAILED build is terminal: drop and fall
                    # through to recreate (the same self-heal `drop()` relies
                    # on). Wait for the drop to clear first -- create_search_index
                    # rejects a name that still exists while the old index is
                    # DELETING.
                    logger.warning(
                        f"[{self.workspace}] vector index {self._index_name} is FAILED "
                        f"and non-queryable; dropping and recreating it"
                    )
                    await self._data.drop_search_index(self._index_name)
                    await self._wait_for_search_index_absent(self._index_name)
                    break

                logger.info(
                    f"[{self.workspace}] vector index {self._index_name} already exists with matching dimensions ({expected_dim})"
                )
                return

            search_index_model = SearchIndexModel(
                definition={
                    "fields": [
                        {
                            "type": "vector",
                            "numDimensions": self.embedding_func.embedding_dim,  # Ensure correct dimensions
                            "path": "vector",
                            "similarity": "cosine",  # Options: euclidean, cosine, dotProduct
                        }
                    ]
                },
                name=self._index_name,
                type="vectorSearch",
            )

            await self._data.create_search_index(search_index_model)
            logger.info(
                f"[{self.workspace}] Vector index {self._index_name} created successfully."
            )

        except PyMongoError as e:
            error_msg = f"[{self.workspace}] Error creating vector index {self._index_name}: {e}"
            logger.error(error_msg)
            raise SystemExit(
                f"Failed to create MongoDB vector index. Program cannot continue. {error_msg}"
            )

    async def upsert(self, data: dict[str, dict[str, Any]]) -> None:
        """Buffer vector docs for embedding and batched flush.

        Embedding deliberately does NOT happen here: repeated upserts of
        the same id, or many small batches, collapse into a single
        flush-time embedding pass. Reads observe pending docs via the
        same lock for read-your-writes.
        """
        if not data:
            return

        current_time = int(time.time())

        pending_docs: list[tuple[str, _PendingVectorDoc]] = []
        for i, (k, v) in enumerate(data.items(), start=1):
            source = {
                "_id": k,
                "created_at": current_time,
                **{k1: v1 for k1, v1 in v.items() if k1 in self.meta_fields},
            }
            pending_docs.append(
                (
                    k,
                    _PendingVectorDoc(source=source, content=v["content"]),
                )
            )
            await _cooperative_yield(i)

        # Installing a fresh _PendingVectorDoc invalidates any vector
        # cached by a prior get_vectors_by_ids() call on a stale revision.
        async with self._flush_lock:
            for doc_id, pdoc in pending_docs:
                self._pending_vector_deletes.discard(doc_id)
                self._pending_vector_docs[doc_id] = pdoc

    async def query(
        self, query: str, top_k: int, query_embedding: list[float] = None
    ) -> list[dict[str, Any]]:
        """Queries the vector database using Atlas Vector Search.

        Reads from the server-side index only; buffered upserts and deletes
        are NOT visible until ``index_done_callback`` / ``finalize`` flushes
        them. Callers that need read-your-writes for a freshly upserted id
        should use ``get_by_id`` / ``get_by_ids`` (which consult the buffer)
        or flush first. Matches the deferred-embedding contract used by
        OpenSearch / FAISS / Nano.
        """
        if query_embedding is not None:
            # Convert numpy array to list if needed for MongoDB compatibility
            if hasattr(query_embedding, "tolist"):
                query_vector = query_embedding.tolist()
            else:
                query_vector = list(query_embedding)
        else:
            # Generate the embedding
            embedding = await self.embedding_func(
                [query], context="query", _priority=DEFAULT_QUERY_PRIORITY
            )  # higher priority for query
            # Convert numpy array to a list to ensure compatibility with MongoDB
            query_vector = embedding[0].tolist()

        # Define the aggregation pipeline with the converted query vector
        pipeline = [
            {
                "$vectorSearch": {
                    "index": self._index_name,  # Use stored index name for consistency
                    "path": "vector",
                    "queryVector": query_vector,
                    "numCandidates": 100,  # Adjust for performance
                    "limit": top_k,
                }
            },
            {"$addFields": {"score": {"$meta": "vectorSearchScore"}}},
            {"$match": {"score": {"$gte": self.cosine_better_than_threshold}}},
            {"$project": {"vector": 0}},
        ]

        # Execute the aggregation pipeline
        cursor = await self._data.aggregate(pipeline, allowDiskUse=True)
        results = await cursor.to_list(length=None)

        # Format and return the results with created_at field
        return [
            {
                **doc,
                "id": doc["_id"],
                "distance": doc.get("score", None),
                "created_at": doc.get("created_at"),  # Include created_at field
            }
            for doc in results
        ]

    async def index_done_callback(self) -> None:
        """Flush buffered vector ops; Mongo persists automatically once written."""
        await self._flush_pending_vector_ops()

    async def drop_pending_index_ops(self) -> None:
        """Discard buffered upserts/deletes (pipeline aborting on error)."""
        async with self._flush_lock:
            self._pending_vector_docs.clear()
            self._pending_vector_deletes.clear()

    async def _flush_pending_vector_ops(self) -> None:
        """Flush buffered vector upserts and deletes in batched bulk writes.

        Embedding runs *inside* this lock (not in `upsert` or lock-free):
        it makes deferred embedding and the bulk write atomic against
        concurrent upserts and destructive mutations. Any failure (embed
        or server write) raises and leaves both buffers intact; the next
        `index_done_callback` retries automatically.

        Concurrency invariant: ``_flush_lock`` is a non-reentrant asyncio
        lock. Callers MUST NOT hold it when invoking this method --
        re-entry would deadlock. The only in-tree callers are
        ``index_done_callback`` and ``finalize``, both lock-free.
        """
        async with self._flush_lock:
            if not self._pending_vector_docs and not self._pending_vector_deletes:
                return
            if self._data is None:
                return

            pending_docs = self._pending_vector_docs
            pending_deletes = self._pending_vector_deletes

            docs_to_embed: list[tuple[str, _PendingVectorDoc]] = [
                (doc_id, pdoc)
                for doc_id, pdoc in pending_docs.items()
                if pdoc.vector is None
            ]

            if docs_to_embed:
                contents = [pdoc.content for _, pdoc in docs_to_embed]
                batches = [
                    contents[i : i + self._max_batch_size]
                    for i in range(0, len(contents), self._max_batch_size)
                ]
                logger.info(
                    f"[{self.workspace}] {self.namespace} flush: embedding "
                    f"{len(docs_to_embed)} vectors in {len(batches)} batch(es) "
                    f"(batch_num={self._max_batch_size})"
                )
                try:
                    embeddings_list = await asyncio.gather(
                        *[
                            self.embedding_func(batch, context="document")
                            for batch in batches
                        ]
                    )
                except Exception as e:
                    logger.error(
                        f"[{self.workspace}] Error embedding pending vector ops "
                        f"(upserts={len(docs_to_embed)}): {e}"
                    )
                    raise

                embeddings = np.concatenate(embeddings_list)
                if len(embeddings) != len(docs_to_embed):
                    raise RuntimeError(
                        f"[{self.workspace}] Embedding count mismatch: expected "
                        f"{len(docs_to_embed)}, got {len(embeddings)}"
                    )
                for i, ((_, pdoc), embedding) in enumerate(
                    zip(docs_to_embed, embeddings), start=1
                ):
                    pdoc.vector = np.array(embedding, dtype=np.float32).tolist()
                    await _cooperative_yield(i)

            # Assemble final upsert payload. After the embed loop above every
            # pending doc has a non-None vector (count-mismatch was checked),
            # so we can iterate without re-guarding. Each full_doc carries its
            # own "_id" (from source), matching the UpdateOne filter key.
            ids_to_commit: list[str] = list(pending_docs.keys())
            list_data: list[dict[str, Any]] = [
                {**pending_docs[doc_id].source, "vector": pending_docs[doc_id].vector}
                for doc_id in ids_to_commit
            ]

            try:
                if list_data:
                    # Split the upsert into batches that stay under the server-side
                    # bulk-command message limit and bound peak memory. Fail-fast:
                    # any batch failure raises immediately and the full buffer is
                    # retained for the next flush (upsert/delete are idempotent).
                    # Logging is kept aligned with MilvusVectorDBStorage; the
                    # batching maths is shared via _chunk_by_budget.
                    upsert_batches = _chunk_by_budget(
                        list_data,
                        _estimate_doc_bytes,
                        self._max_upsert_payload_bytes,
                        self._max_upsert_records_per_batch,
                    )
                    if len(upsert_batches) > 1:
                        logger.info(
                            f"[{self.workspace}] {self.namespace} flush: upsert split into "
                            f"{len(upsert_batches)} batches for {len(list_data)} records "
                            f"(max_payload={self._max_upsert_payload_bytes} batch={self._max_upsert_records_per_batch})"
                        )
                    for batch_index, (records_batch, estimated_bytes) in enumerate(
                        upsert_batches, 1
                    ):
                        if (
                            len(records_batch) == 1
                            and self._max_upsert_payload_bytes > 0
                            and estimated_bytes > self._max_upsert_payload_bytes
                        ):
                            logger.warning(
                                f"[{self.workspace}] {self.namespace} flush: single record "
                                f"id={records_batch[0].get('_id')} estimated {estimated_bytes} bytes "
                                f"exceeds {self._max_upsert_payload_bytes}"
                            )
                        logger.debug(
                            f"[{self.workspace}] MongoDB upsert batch {batch_index}/{len(upsert_batches)}: "
                            f"records={len(records_batch)}, estimated_payload_bytes={estimated_bytes}"
                        )
                        await self._data.bulk_write(
                            [
                                UpdateOne(
                                    {"_id": doc["_id"]}, {"$set": doc}, upsert=True
                                )
                                for doc in records_batch
                            ],
                            ordered=False,
                        )
                if pending_deletes:
                    # Chunk deletes by record count; _ids are short strings so a
                    # count cap is enough to stay under the bulk message limit.
                    # delete_many($in) is the 1:1 equivalent of a batched delete.
                    delete_ids = list(pending_deletes)
                    delete_chunk = (
                        self._max_delete_records_per_batch
                        if self._max_delete_records_per_batch > 0
                        else len(delete_ids)
                    )
                    for i in range(0, len(delete_ids), delete_chunk):
                        await self._data.delete_many(
                            {"_id": {"$in": delete_ids[i : i + delete_chunk]}}
                        )
            except Exception as e:
                logger.error(
                    f"[{self.workspace}] Error flushing vector ops "
                    f"(upserts={len(pending_docs)}, "
                    f"deletes={len(pending_deletes)}): {e}"
                )
                raise

            # On success, clear the buffers in-place so external references
            # (e.g. drop()) see the cleared state.
            for doc_id in ids_to_commit:
                pending_docs.pop(doc_id, None)
            pending_deletes.clear()

    async def delete(self, ids: list[str]) -> None:
        """Buffer vector deletes for batched flush."""
        if not ids:
            return
        if isinstance(ids, set):
            ids = list(ids)
        async with self._flush_lock:
            for doc_id in ids:
                self._pending_vector_docs.pop(doc_id, None)
                self._pending_vector_deletes.add(doc_id)
        logger.debug(
            f"[{self.workspace}] Buffered delete for {len(ids)} vectors in {self.namespace}"
        )

    async def delete_entity(self, entity_name: str) -> None:
        """Buffer an entity vector delete by computing its hash ID."""
        entity_id = compute_mdhash_id(entity_name, prefix="ent-")
        async with self._flush_lock:
            self._pending_vector_docs.pop(entity_id, None)
            self._pending_vector_deletes.add(entity_id)
        logger.debug(
            f"[{self.workspace}] Buffered delete for entity {entity_name} (id={entity_id})"
        )

    async def delete_entity_relation(self, entity_name: str) -> None:
        """Delete all relation vectors where entity appears as src or tgt.

        The whole method runs under ``_flush_lock`` so the server-side find
        + delete cannot interleave with an in-flight bulk write. Server-side
        failures are re-raised (no log-and-swallow): the caller decides
        whether to retry.

        Buffer semantics — post-prune with caller short-circuit contract:
            Matching pending upserts in ``_pending_vector_docs`` are
            pruned **only after** the server-side ``delete_many``
            succeeds. On failure the pending buffer stays intact and
            the exception propagates so the caller (``adelete_by_entity``
            in ``utils_graph.py``) can short-circuit before
            ``_persist_graph_updates`` flushes a half-cleaned buffer.
        """

        def _prune_pending() -> None:
            for doc_id in [
                k
                for k, v in self._pending_vector_docs.items()
                if v.source.get("src_id") == entity_name
                or v.source.get("tgt_id") == entity_name
            ]:
                self._pending_vector_docs.pop(doc_id, None)

        async with self._flush_lock:
            if self._data is None:
                # No server state to mutate; buffer prune is the only
                # delete intent we can record.
                _prune_pending()
                return

            # _id is the only field we need from the find; project to keep
            # the cursor light.
            relations_cursor = self._data.find(
                {"$or": [{"src_id": entity_name}, {"tgt_id": entity_name}]},
                {"_id": 1},
            )
            relations = await relations_cursor.to_list(length=None)

            if not relations:
                # No server rows to delete — still safe to prune any
                # pending upserts so they can't re-create the relation.
                _prune_pending()
                logger.debug(
                    f"[{self.workspace}] No relations found for entity {entity_name}"
                )
                return

            relation_ids = [relation["_id"] for relation in relations]
            await self._data.delete_many({"_id": {"$in": relation_ids}})
            # Server-side delete succeeded — safe to prune the pending
            # buffer so subsequent flushes don't re-upsert the deleted
            # relations.
            _prune_pending()
            logger.debug(
                f"[{self.workspace}] Deleted {len(relation_ids)} relations for {entity_name}"
            )

    async def get_by_id(self, id: str) -> dict[str, Any] | None:
        """Get vector data by its ID, with read-your-writes against the buffer.

        Pending buffer hits never include the `vector` field; server-side
        fallback projects it out for parity.
        """
        async with self._flush_lock:
            if id in self._pending_vector_deletes:
                return None
            pending = self._pending_vector_docs.get(id)
            if pending is not None:
                doc = dict(pending.source)
                # Surface both _id (Mongo native) and id (API expectation).
                doc.setdefault("_id", id)
                doc["id"] = id
                return doc

        try:
            result = await self._data.find_one({"_id": id}, {"vector": 0})
            if result:
                result_dict = dict(result)
                if "_id" in result_dict and "id" not in result_dict:
                    result_dict["id"] = result_dict["_id"]
                return result_dict
            return None
        except Exception as e:
            logger.error(
                f"[{self.workspace}] Error retrieving vector data for ID {id}: {e}"
            )
            return None

    async def get_by_ids(self, ids: list[str]) -> list[dict[str, Any]]:
        """Get multiple vector data by their IDs (read-your-writes), preserving order."""
        if not ids:
            return []

        buffered: dict[str, dict[str, Any] | None] = {}
        remaining: list[str] = []
        async with self._flush_lock:
            for doc_id in ids:
                if doc_id in self._pending_vector_deletes:
                    buffered[doc_id] = None
                    continue
                pending = self._pending_vector_docs.get(doc_id)
                if pending is not None:
                    doc = dict(pending.source)
                    doc.setdefault("_id", doc_id)
                    doc["id"] = doc_id
                    buffered[doc_id] = doc
                    continue
                remaining.append(doc_id)

        formatted_map: dict[str, dict[str, Any]] = {}
        if remaining:
            try:
                cursor = self._data.find({"_id": {"$in": remaining}}, {"vector": 0})
                results = await cursor.to_list(length=None)
                for result in results:
                    result_dict = dict(result)
                    if "_id" in result_dict and "id" not in result_dict:
                        result_dict["id"] = result_dict["_id"]
                    key = str(result_dict.get("id", result_dict.get("_id")))
                    formatted_map[key] = result_dict
            except Exception as e:
                logger.error(
                    f"[{self.workspace}] Error retrieving vector data for IDs {remaining}: {e}"
                )
                return []

        return [
            buffered[doc_id] if doc_id in buffered else formatted_map.get(str(doc_id))
            for doc_id in ids
        ]

    async def get_vectors_by_ids(self, ids: list[str]) -> dict[str, list[float]]:
        """Get vector embeddings for given IDs, with read-your-writes.

        Pending docs whose vector hasn't been embedded yet are embedded
        lazily inside the lock; the resulting vector is cached on the
        buffered `_PendingVectorDoc` so the next flush won't re-embed.

        Visibility caveat for ids not in the buffer: the server-side
        ``find`` fallback runs *outside* ``_flush_lock``. A concurrent
        ``delete()`` that lands between lock release and the cursor
        read only buffers the delete -- the old vector is still on disk
        until the next flush, so this method may return a stale vector
        for an id that has been buffered for deletion. This is
        best-effort read-after-uncommitted-delete and matches the
        ``query()`` contract: callers needing strict consistency must
        ``index_done_callback()`` first.
        """
        if not ids:
            return {}

        result: dict[str, list[float]] = {}
        remaining: list[str] = []
        async with self._flush_lock:
            docs_to_embed: list[tuple[str, _PendingVectorDoc]] = []
            for doc_id in ids:
                if doc_id in self._pending_vector_deletes:
                    continue
                pending = self._pending_vector_docs.get(doc_id)
                if pending is not None:
                    if pending.vector is None:
                        docs_to_embed.append((doc_id, pending))
                    else:
                        result[doc_id] = pending.vector
                    continue
                remaining.append(doc_id)

            if docs_to_embed:
                contents = [pdoc.content for _, pdoc in docs_to_embed]
                batches = [
                    contents[i : i + self._max_batch_size]
                    for i in range(0, len(contents), self._max_batch_size)
                ]
                try:
                    embeddings_list = await asyncio.gather(
                        *[
                            self.embedding_func(batch, context="document")
                            for batch in batches
                        ]
                    )
                except Exception as e:
                    logger.error(
                        f"[{self.workspace}] Error lazily embedding pending vectors "
                        f"(upserts={len(docs_to_embed)}): {e}"
                    )
                    raise
                embeddings = np.concatenate(embeddings_list)
                if len(embeddings) != len(docs_to_embed):
                    raise RuntimeError(
                        f"[{self.workspace}] Embedding count mismatch: expected "
                        f"{len(docs_to_embed)}, got {len(embeddings)}"
                    )
                for i, ((doc_id, pdoc), embedding) in enumerate(
                    zip(docs_to_embed, embeddings), start=1
                ):
                    pdoc.vector = np.array(embedding, dtype=np.float32).tolist()
                    result[doc_id] = pdoc.vector
                    await _cooperative_yield(i)

        if not remaining:
            return result

        try:
            cursor = self._data.find(
                {"_id": {"$in": remaining}}, {"_id": 1, "vector": 1}
            )
            results = await cursor.to_list(length=None)
            for row in results:
                if row and "vector" in row and "_id" in row:
                    result[row["_id"]] = row["vector"]
            return result
        except PyMongoError as e:
            logger.error(f"[{self.workspace}] Error getting vectors: {e}")
            return result

    async def drop(self) -> dict[str, str]:
        """Drop all documents and recreate the vector index. Destructive.

        MUST only be called when ``pipeline_status`` is idle (see the
        Pipeline concurrency contract in ``AGENTS.md``); the only
        in-tree caller ``clear_documents`` enforces this.

        Caveat — only this instance's buffers are cleared. Other
        ``MongoVectorDBStorage`` instances aliased onto the same
        ``final_namespace`` (multi-worker processes, or distinct
        workspaces collapsed by ``MONGODB_WORKSPACE``) keep their own
        buffers; a sibling whose prior flush failed and left buffers
        intact will, on its next flush, bulk-write those stale rows into
        the freshly recreated collection. Direct callers bypassing the
        idle precondition MUST flush every aliased instance first.

        Returns:
            dict[str, str]: ``{"status": "success"|"error", "message": str}``
        """
        try:
            async with self._flush_lock:
                # Discard any buffered writes before the collection is wiped;
                # a concurrent flush would otherwise resurrect them.
                self._pending_vector_docs.clear()
                self._pending_vector_deletes.clear()

                # Delete all documents
                result = await self._data.delete_many({})
                deleted_count = result.deleted_count

                # Recreate vector index
                await self.create_vector_index_if_not_exists()

            logger.info(
                f"[{self.workspace}] Dropped {deleted_count} documents from vector storage {self._collection_name} and recreated vector index"
            )
            return {
                "status": "success",
                "message": f"{deleted_count} documents dropped and vector index recreated",
            }
        except PyMongoError as e:
            logger.error(
                f"[{self.workspace}] Error dropping vector storage {self._collection_name}: {e}"
            )
            return {"status": "error", "message": str(e)}


async def get_or_create_collection(db: AsyncDatabase, collection_name: str):
    collection_names = await db.list_collection_names()

    if collection_name not in collection_names:
        collection = await db.create_collection(collection_name)
        logger.info(f"Created collection: {collection_name}")
        return collection
    else:
        logger.debug(f"Collection '{collection_name}' already exists.")
        return db.get_collection(collection_name)
