embeddings/database.py

"""Qdrant database manager for hybrid (dense + sparse) vector storage."""

from __future__ import annotations

from typing import Any

from qdrant_client import QdrantClient, models
from qdrant_client.models import (
    Distance,
    PointStruct,
    SparseVector,
    SparseVectorParams,
    VectorParams,
)

from src.config import (
    EMBEDDING_DIMENSION,
    QDRANT_API_KEY,
    QDRANT_COLLECTION_NAME,
    QDRANT_TIMEOUT,
    QDRANT_URL,
)

# Batch size for upsert operations
_UPSERT_BATCH_SIZE = 64


def _build_client() -> QdrantClient:
    """Create a Qdrant client, preferring cloud when credentials are available."""
    if QDRANT_URL and QDRANT_API_KEY and "your_" not in QDRANT_URL:
        print(f"[QDRANT] Connecting to cloud: {QDRANT_URL} (timeout={QDRANT_TIMEOUT}s)")
        return QdrantClient(
            url=QDRANT_URL,
            api_key=QDRANT_API_KEY,
            timeout=QDRANT_TIMEOUT,
        )

    # Fallback: local persistent storage
    local_path = "data/qdrant_local"
    print(f"[QDRANT] Using local storage: {local_path}")
    return QdrantClient(path=local_path, timeout=QDRANT_TIMEOUT)


class QdrantManager:
    """Manage a Qdrant collection that stores both dense and sparse vectors."""

    def __init__(
        self,
        collection_name: str = QDRANT_COLLECTION_NAME,
        client: QdrantClient | None = None,
    ) -> None:
        self.collection_name = collection_name
        self.client = client or _build_client()

    # ------------------------------------------------------------------
    # Collection lifecycle
    # ------------------------------------------------------------------

    def init_collection(self, *, recreate: bool = False) -> None:
        """Create the collection if it doesn't exist.

        If *recreate* is ``True``, delete and recreate it.
        """
        exists = self.client.collection_exists(self.collection_name)

        if exists and not recreate:
            print(f"[QDRANT] Collection '{self.collection_name}' already exists — skipping creation.")
            return

        if exists and recreate:
            self.client.delete_collection(self.collection_name)
            print(f"[QDRANT] Deleted existing collection '{self.collection_name}'.")

        self.client.create_collection(
            collection_name=self.collection_name,
            vectors_config={
                "dense": VectorParams(
                    size=EMBEDDING_DIMENSION,
                    distance=Distance.COSINE,
                ),
            },
            sparse_vectors_config={
                "sparse": SparseVectorParams(),
            },
        )
        print(f"[QDRANT] Created collection '{self.collection_name}' (dense={EMBEDDING_DIMENSION}d + sparse).")

    # ------------------------------------------------------------------
    # Upsert
    # ------------------------------------------------------------------

    def upsert_chunks(
        self,
        *,
        chunk_ids: list[str],
        contents: list[str],
        metadatas: list[dict[str, Any]],
        dense_vectors: list[list[float]],
        sparse_vectors: list[dict[str, Any]],
    ) -> int:
        """Upload chunks with both dense and sparse vectors to Qdrant.

        Returns the number of points upserted.
        """
        n = len(chunk_ids)
        if not (n == len(contents) == len(metadatas) == len(dense_vectors) == len(sparse_vectors)):
            raise ValueError("All input lists must have the same length.")

        total_upserted = 0

        for start in range(0, n, _UPSERT_BATCH_SIZE):
            end = min(start + _UPSERT_BATCH_SIZE, n)
            points: list[PointStruct] = []

            for i in range(start, end):
                # Qdrant requires integer or UUID point IDs.
                # We use a deterministic hash of the chunk_id string.
                point_id = _stable_id(chunk_ids[i])

                payload = {
                    "chunk_id": chunk_ids[i],
                    "content": contents[i],
                    **metadatas[i],
                }

                sv = sparse_vectors[i]
                point = PointStruct(
                    id=point_id,
                    vector={
                        "dense": dense_vectors[i],
                        "sparse": SparseVector(
                            indices=sv["indices"],
                            values=sv["values"],
                        ),
                    },
                    payload=payload,
                )
                points.append(point)

            self.client.upsert(collection_name=self.collection_name, points=points)
            total_upserted += len(points)
            print(f"[QDRANT] Upserted batch {start}–{end} ({len(points)} points)")

        return total_upserted

    # ------------------------------------------------------------------
    # Search
    # ------------------------------------------------------------------

    def dense_search(
        self,
        *,
        dense_vector: list[float],
        limit: int = 5,
    ) -> list[models.ScoredPoint]:
        """Search by dense vector only (cosine similarity)."""
        results = self.client.query_points(
            collection_name=self.collection_name,
            query=dense_vector,
            using="dense",
            limit=limit,
        )
        return results.points

    def sparse_search(
        self,
        *,
        sparse_vector: dict[str, Any],
        limit: int = 5,
    ) -> list[models.ScoredPoint]:
        """Search by sparse (BM25) vector only."""
        sv = SparseVector(
            indices=sparse_vector["indices"],
            values=sparse_vector["values"],
        )
        results = self.client.query_points(
            collection_name=self.collection_name,
            query=sv,
            using="sparse",
            limit=limit,
        )
        return results.points

    def hybrid_search(
        self,
        *,
        dense_vector: list[float],
        sparse_vector: dict[str, Any],
        limit: int = 5,
        dense_limit: int | None = None,
        sparse_limit: int | None = None,
        fusion_limit: int | None = None,
    ) -> list[models.ScoredPoint]:
        """Run a hybrid query using Qdrant's built-in RRF fusion."""
        sv = SparseVector(
            indices=sparse_vector["indices"],
            values=sparse_vector["values"],
        )
        dense_prefetch_limit = dense_limit or limit
        sparse_prefetch_limit = sparse_limit or limit
        final_limit = fusion_limit or limit

        results = self.client.query_points(
            collection_name=self.collection_name,
            prefetch=[
                models.Prefetch(
                    query=dense_vector,
                    using="dense",
                    limit=dense_prefetch_limit,
                ),
                models.Prefetch(
                    query=sv,
                    using="sparse",
                    limit=sparse_prefetch_limit,
                ),
            ],
            query=models.FusionQuery(fusion=models.Fusion.RRF),
            limit=final_limit,
        )
        return results.points

    # ------------------------------------------------------------------
    # Info
    # ------------------------------------------------------------------

    def delete_by_doc_id(self, doc_id: str) -> None:
        """Delete all points whose payload doc_id matches."""
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=models.FilterSelector(
                filter=models.Filter(
                    must=[
                        models.FieldCondition(
                            key="doc_id",
                            match=models.MatchValue(value=doc_id),
                        )
                    ]
                )
            ),
        )

    def count(self) -> int:
        """Return the number of points in the collection."""
        info = self.client.get_collection(self.collection_name)
        return int(info.points_count or 0)


def _stable_id(chunk_id: str) -> int:
    """Produce a positive 64-bit integer from a chunk_id string."""
    import hashlib

    digest = hashlib.sha256(chunk_id.encode()).hexdigest()
    return int(digest[:16], 16)