from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generic, Literal, Sequence, TypeVar

from .chunk import RetrievedChunk
from .document import Document

TDocument = TypeVar("TDocument", bound=Document, covariant=True)


@dataclass(frozen=True)
class WriteResult(Generic[TDocument]):
    action: Literal["inserted", "replaced", "skipped"]
    document: TDocument
    replaced_document: TDocument | None = None


class BaseStore(ABC):
    """Abstract base class for vector stores.

    A store is responsible for storing documents and their embeddings,
    and retrieving relevant chunks based on similarity search.

    Subclasses must implement all abstract methods to provide a concrete
    storage backend:

    - :py:class:`raghilda.store.DuckDBStore`: local storage with embedding and BM25 search.
    - :py:class:`raghilda.store.ChromaDBStore`: local storage using ChromaDB.
    - :py:class:`raghilda.store.OpenAIStore`: hosted storage using OpenAI's Vector Store API.
    """

    @staticmethod
    @abstractmethod
    def connect(*args, **kwargs) -> "BaseStore":
        """Connect to an existing store.

        Returns
        -------
        BaseStore
            A connected store instance.
        """
        pass

    @staticmethod
    @abstractmethod
    def create(*args, **kwargs) -> "BaseStore":
        """Create a new store.

        Returns
        -------
        BaseStore
            A newly created store instance.
        """
        pass

    @abstractmethod
    def upsert(
        self,
        document: Document,
        *,
        skip_if_unchanged: bool = True,
    ) -> WriteResult[Document]:
        """Upsert a document into the store.

        Insert or replace a document in the store.

        Parameters
        ----------
        document
            The document to upsert.
        skip_if_unchanged
            If True (default), skip the write when the existing document
            for the same identity key already has identical content and
            chunk metadata. This helps avoid unnecessary embedding work.
        """
        pass

    @abstractmethod
    def retrieve(
        self, text: str, top_k: int, *args, **kwargs
    ) -> Sequence[RetrievedChunk]:
        """Retrieve the most similar chunks to the given text.

        Parameters
        ----------
        text
            The query text to search for.
        top_k
            The maximum number of chunks to return.

        Returns
        -------
        Sequence[RetrievedChunk]
            The most similar chunks, ordered by relevance.
        """
        pass

    @abstractmethod
    def size(self) -> int:
        """Count the number of documents in the store.

        Returns
        -------
        int
            The number of documents (not chunks) in the store.
        """
        pass
