"""LlamaIndex vector store adapter for RAGU's BaseVectorStorage interface."""

from typing import List

try:
    from llama_index.core.vector_stores.types import (
        BasePydanticVectorStore,
        VectorStoreQuery,
        VectorStoreQueryMode,
    )
    from llama_index.core.schema import TextNode
except ImportError:
    raise ImportError(
        "Please install 'llama-index-core' for LlamaIndex integration:\n"
        "    pip install llama-index-core\n"
        "Or install RAGU with llama-index support: pip install graph_ragu[llama-index]"
    )

from ragu.storage.base_storage import BaseVectorStorage
from ragu.storage.types import Embedding, EmbeddingHit


class LlamaIndexVectorStorage(BaseVectorStorage):
    """Adapter that wraps any LlamaIndex ``BasePydanticVectorStore`` behind RAGU's
    :class:`~ragu.storage.base_storage.BaseVectorStorage` interface.

    The user constructs a LlamaIndex vector store (Qdrant, Faiss, Chroma, etc.)
    and passes it to this adapter.  The adapter translates between the two APIs.

    :param llama_index_store: An already-constructed LlamaIndex vector store.
    :type llama_index_store: BasePydanticVectorStore
    :param embedding_dim: Dimensionality of the embedding vectors.
    :type embedding_dim: int
    :param kwargs: Additional keyword arguments forwarded to the base class.
    """

    def __init__(
        self,
        llama_index_store: BasePydanticVectorStore,
        embedding_dim: int,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._store = llama_index_store
        self.embedding_dim = embedding_dim

    async def upsert(self, data: List[Embedding]) -> None:
        """Insert or update embeddings by converting them to LlamaIndex ``TextNode`` objects.

        :param data: Embedding records to upsert.
        :type data: List[Embedding]
        """
        if not data:
            return

        nodes: List[TextNode] = []
        for emb in data:
            if emb.vector is None:
                continue
            vector = list(emb.vector) if not isinstance(emb.vector, list) else emb.vector
            node = TextNode(
                id_=emb.id,
                text="",
                embedding=vector,
                metadata=emb.metadata,
            )
            nodes.append(node)

        if not nodes:
            return

        await self._store.async_add(nodes)

    async def query(self, vector: Embedding, top_k: int = 5) -> List[EmbeddingHit]:
        """Query for nearest neighbors.

        :param vector: Query embedding.
        :type vector: Embedding
        :param top_k: Number of nearest neighbors to return.
        :type top_k: int
        :returns: List of matched records with distance scores.
        :rtype: List[EmbeddingHit]
        """
        query_vector = list(vector.vector) if not isinstance(vector.vector, list) else vector.vector

        vstore_query = VectorStoreQuery(
            query_embedding=query_vector,
            similarity_top_k=top_k,
            mode=VectorStoreQueryMode.DEFAULT,
        )

        result = await self._store.aquery(vstore_query)

        hits: List[EmbeddingHit] = []

        if result.nodes is not None:
            for i, node in enumerate(result.nodes):
                distance = result.similarities[i] if result.similarities else 0.0
                hits.append(
                    EmbeddingHit(
                        id=node.node_id,
                        distance=float(distance),
                        metadata=node.metadata,
                    )
                )
        elif result.ids is not None:
            for i, node_id in enumerate(result.ids):
                distance = result.similarities[i] if result.similarities else 0.0
                hits.append(
                    EmbeddingHit(
                        id=node_id,
                        distance=float(distance),
                        metadata={},
                    )
                )

        return hits

    async def delete(self, ids: List[str]) -> None:
        """Delete records by IDs.

        :param ids: Record identifiers to remove.
        :type ids: List[str]
        """
        if not ids:
            return
        await self._store.adelete_nodes(ids)

    async def index_start_callback(self):
        """Pre-index hook (no-op)."""
        pass

    async def index_done_callback(self):
        """Post-index hook. Calls persist on the underlying store if supported."""
        if hasattr(self._store, "persist"):
            self._store.persist()

    async def query_done_callback(self):
        """Post-query hook (no-op)."""
        pass
