library-desk/src/services/vector_service.py

"""
Vector service for Library Desk Qdrant operations.

Handles semantic search, document chunking, and embeddings.
"""

import re
import time
import hashlib
import uuid
from typing import List, Dict, Any, Optional
import logging

from src.clients.qdrant_client import QdrantClientWrapper
from src.clients.wikijs_client import WikiJSClient
from src.clients.ollama_client import OllamaClient
from src.core.multi_tenancy import get_qdrant_collection_name, is_path_in_user_namespace
from src.models.vector import (
    SearchResult, SearchResponse, VectorUpdateSummary,
    DocumentChunk, CollectionInfo, CollectionListResponse
)

logger = logging.getLogger(__name__)


class VectorService:
    """
    Service for Qdrant vector operations.

    Responsibilities:
    - Document chunking
    - Embedding generation
    - Semantic search
    - Vector CRUD operations
    """

    def __init__(
        self,
        qdrant_client: QdrantClientWrapper,
        wikijs_client: WikiJSClient,
        ollama_client: OllamaClient,
        chunk_size: int = 500,
        chunk_overlap: int = 50
    ):
        """
        Initialize vector service.

        Args:
            qdrant_client: Qdrant database client
            wikijs_client: Wiki.js client for fetching pages
            ollama_client: Ollama client for embeddings
            chunk_size: Target chunk size in tokens (approximate)
            chunk_overlap: Overlap between chunks in tokens
        """
        self.qdrant = qdrant_client
        self.wiki = wikijs_client
        self.ollama = ollama_client
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def _chunk_text(self, text: str) -> List[str]:
        """
        Chunk text into overlapping segments.

        Simple word-based chunking for now.
        TODO: Use tiktoken or similar for token-accurate chunking.

        Args:
            text: Text to chunk

        Returns:
            List of text chunks
        """
        # Remove extra whitespace
        text = re.sub(r'\s+', ' ', text).strip()

        # Split into words (approximates tokens)
        words = text.split()

        if len(words) <= self.chunk_size:
            return [text]

        chunks = []
        start = 0

        while start < len(words):
            end = start + self.chunk_size
            chunk_words = words[start:end]
            chunks.append(' '.join(chunk_words))

            # Move start forward with overlap
            start = end - self.chunk_overlap

        return chunks

    async def update_from_page(
        self,
        page_id: int,
        user: str,
        force_refresh: bool = False
    ) -> VectorUpdateSummary:
        """
        Update vector embeddings from a wiki page.

        Chunks the page content, generates embeddings, and upserts to Qdrant.

        Args:
            page_id: Wiki page ID
            user: User identifier
            force_refresh: Force re-embedding even if unchanged

        Returns:
            Summary of update operation
        """
        start_time = time.time()

        try:
            # Fetch page from Wiki.js
            page = await self.wiki.get_page(page_id)
            if not page:
                raise ValueError(f"Page {page_id} not found")

            # TENANT ISOLATION: only pages inside the user's own wiki
            # namespace may be embedded into that user's collection.
            # Without this check any tenant could ingest (and then read)
            # another tenant's wiki content.
            if not is_path_in_user_namespace(page.get("path", ""), user):
                raise ValueError(
                    f"Page {page_id} (path: {page.get('path')!r}) is outside "
                    f"user '{user}' namespace - refusing cross-tenant ingestion"
                )

            # Get collection name for user
            collection_name = get_qdrant_collection_name(user)

            # Ensure collection exists
            await self.qdrant.ensure_collection(collection_name)

            # Extract content
            content = page.get("content", "")
            title = page.get("title", "")
            path = page.get("path", "")

            if not content:
                logger.warning(f"Page {page_id} has no content, skipping vector update")
                return VectorUpdateSummary(
                    page_id=page_id,
                    page_title=title,
                    processing_time_ms=(time.time() - start_time) * 1000,
                    success=True
                )

            # Chunk the content
            chunks = self._chunk_text(content)
            logger.info(f"Split page {page_id} into {len(chunks)} chunks")

            # Delete existing chunks for this page
            deleted_count = await self.qdrant.delete_by_filter(
                collection_name=collection_name,
                filter_conditions={"page_id": page_id}
            )

            # Generate embeddings and upsert chunks
            chunks_created = 0
            for idx, chunk_text in enumerate(chunks):
                # Generate deterministic UUID from page_id and chunk_index
                chunk_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, f"page_{page_id}_chunk_{idx}"))

                # Generate embedding
                embedding = await self.ollama.embed(chunk_text)
                if not embedding:
                    logger.error(f"Failed to generate embedding for chunk {chunk_id}")
                    continue

                # Prepare metadata
                metadata = {
                    "page_id": page_id,
                    "page_title": title,
                    "page_path": path,
                    "chunk_index": idx,
                    "chunk_text": chunk_text,
                    "user": user
                }

                # Upsert to Qdrant
                success = await self.qdrant.upsert_vector(
                    collection_name=collection_name,
                    vector_id=chunk_id,
                    vector=embedding,
                    payload=metadata
                )

                if success:
                    chunks_created += 1

            processing_time_ms = (time.time() - start_time) * 1000

            logger.info(
                f"Updated vectors for page {page_id}: "
                f"{chunks_created} chunks created, {deleted_count} old chunks deleted"
            )

            return VectorUpdateSummary(
                page_id=page_id,
                page_title=title,
                chunks_created=chunks_created,
                chunks_deleted=deleted_count,
                total_chunks=chunks_created,
                embedding_dim=len(embedding) if embedding else 768,
                processing_time_ms=processing_time_ms,
                success=True
            )

        except Exception as e:
            processing_time_ms = (time.time() - start_time) * 1000
            logger.error(f"Failed to update vectors for page {page_id}: {e}", exc_info=True)

            return VectorUpdateSummary(
                page_id=page_id,
                page_title="Unknown",
                processing_time_ms=processing_time_ms,
                success=False,
                error_message=str(e)
            )

    async def search(
        self,
        query: str,
        user: str,
        limit: int = 10,
        score_threshold: float = 0.5
    ) -> SearchResponse:
        """
        Perform semantic search across user's documents.

        Args:
            query: Search query text
            user: User identifier
            limit: Maximum results to return
            score_threshold: Minimum similarity score (0-1)

        Returns:
            Search results with similarity scores
        """
        start_time = time.time()

        try:
            # Get collection name
            collection_name = get_qdrant_collection_name(user)

            # Check if collection exists
            exists = await self.qdrant.collection_exists(collection_name)
            if not exists:
                logger.info(f"Collection {collection_name} doesn't exist, returning empty results")
                return SearchResponse(
                    query=query,
                    results=[],
                    total=0,
                    user=user
                )

            # Generate query embedding
            query_embedding = await self.ollama.embed(query)
            if not query_embedding:
                raise ValueError("Failed to generate query embedding")

            # Search in Qdrant
            search_results = await self.qdrant.search_vectors(
                collection_name=collection_name,
                query_vector=query_embedding,
                limit=limit,
                score_threshold=score_threshold
            )

            # Convert to SearchResult models
            results = []
            for result in search_results:
                payload = result.get("payload", {})
                results.append(SearchResult(
                    chunk_id=result["id"],
                    page_id=payload.get("page_id", 0),
                    page_title=payload.get("page_title"),
                    page_path=payload.get("page_path"),
                    chunk_index=payload.get("chunk_index", 0),
                    content=payload.get("chunk_text", ""),
                    score=result["score"],
                    metadata=payload
                ))

            query_time_ms = (time.time() - start_time) * 1000
            logger.info(f"Semantic search completed in {query_time_ms:.2f}ms: {len(results)} results")

            return SearchResponse(
                query=query,
                results=results,
                total=len(results),
                user=user
            )

        except Exception as e:
            logger.error(f"Semantic search failed: {e}", exc_info=True)
            return SearchResponse(
                query=query,
                results=[],
                total=0,
                user=user
            )

    async def delete_page_chunks(
        self,
        page_id: int,
        user: str
    ) -> int:
        """
        Delete all chunks for a wiki page.

        Args:
            page_id: Wiki page ID
            user: User identifier

        Returns:
            Number of chunks deleted
        """
        collection_name = get_qdrant_collection_name(user)

        try:
            deleted_count = await self.qdrant.delete_by_filter(
                collection_name=collection_name,
                filter_conditions={"page_id": page_id}
            )

            logger.info(f"Deleted {deleted_count} chunks for page {page_id}")
            return deleted_count

        except Exception as e:
            logger.error(f"Failed to delete chunks for page {page_id}: {e}", exc_info=True)
            return 0

    async def list_collections(self) -> CollectionListResponse:
        """
        List all Qdrant collections.

        Returns:
            List of collections with stats
        """
        try:
            collections_data = await self.qdrant.list_collections()

            collections = []
            for coll in collections_data:
                collections.append(CollectionInfo(
                    name=coll["name"],
                    vectors_count=coll.get("vectors_count", 0),
                    points_count=coll.get("points_count", 0),
                    segments_count=coll.get("segments_count", 0)
                ))

            return CollectionListResponse(
                collections=collections,
                total=len(collections)
            )

        except Exception as e:
            logger.error(f"Failed to list collections: {e}", exc_info=True)
            return CollectionListResponse(
                collections=[],
                total=0
            )

    # ========== Cleanup Methods ==========

    async def delete_document_chunks(
        self,
        document_id: str,
        user: str
    ) -> int:
        """
        Delete all chunks for a document (Document Store).

        Args:
            document_id: Document UUID
            user: User identifier

        Returns:
            Number of chunks deleted
        """
        collection_name = get_qdrant_collection_name(user)

        try:
            deleted_count = await self.qdrant.delete_by_filter(
                collection_name=collection_name,
                filter_conditions={"document_id": document_id}
            )

            logger.info(f"Deleted chunks for document {document_id}")
            return deleted_count

        except Exception as e:
            logger.error(f"Failed to delete chunks for document {document_id}: {e}", exc_info=True)
            return 0

    async def delete_paperless_document_chunks(
        self,
        paperless_id: int,
        user: str
    ) -> int:
        """
        Delete all chunks for a Paperless document.

        Args:
            paperless_id: Paperless-ngx document ID
            user: User identifier

        Returns:
            Number of chunks deleted
        """
        collection_name = get_qdrant_collection_name(user)

        try:
            deleted_count = await self.qdrant.delete_by_filter(
                collection_name=collection_name,
                filter_conditions={
                    "doc_type": "document",
                    "paperless_id": paperless_id
                }
            )

            logger.info(f"Deleted chunks for Paperless document {paperless_id}")
            return deleted_count

        except Exception as e:
            logger.error(f"Failed to delete chunks for Paperless document {paperless_id}: {e}", exc_info=True)
            return 0

    async def delete_collection_chunks(
        self,
        collection_id: str,
        user: str
    ) -> int:
        """
        Delete all chunks for a document collection.

        Args:
            collection_id: Collection UUID
            user: User identifier

        Returns:
            Number of chunks deleted
        """
        collection_name = get_qdrant_collection_name(user)

        try:
            deleted_count = await self.qdrant.delete_by_filter(
                collection_name=collection_name,
                filter_conditions={"collection_id": collection_id}
            )

            logger.info(f"Deleted chunks for collection {collection_id}")
            return deleted_count

        except Exception as e:
            logger.error(f"Failed to delete chunks for collection {collection_id}: {e}", exc_info=True)
            return 0

    async def get_all_chunk_references(
        self,
        user: str
    ) -> List[Dict[str, Any]]:
        """
        Get all chunk references for orphan detection.

        Returns list of {id, page_id, document_id} for all chunks.

        Args:
            user: User identifier

        Returns:
            List of chunk references
        """
        collection_name = get_qdrant_collection_name(user)

        try:
            # Check if collection exists
            exists = await self.qdrant.collection_exists(collection_name)
            if not exists:
                return []

            all_points = await self.qdrant.scroll_all_points(
                collection_name=collection_name,
                batch_size=100,
                with_payload=True
            )

            references = []
            for point in all_points:
                payload = point.get("payload", {})
                references.append({
                    "chunk_id": point["id"],
                    "page_id": payload.get("page_id"),
                    "document_id": payload.get("document_id"),
                    "paperless_id": payload.get("paperless_id"),  # For Paperless documents
                    "collection_id": payload.get("collection_id"),
                    "doc_type": payload.get("doc_type", "wiki")
                })

            logger.info(f"Found {len(references)} chunks for user {user}")
            return references

        except Exception as e:
            logger.error(f"Failed to get chunk references: {e}", exc_info=True)
            return []

    async def purge_chunks_by_ids(
        self,
        user: str,
        chunk_ids: List[str]
    ) -> int:
        """
        Delete specific chunks by their IDs.

        Args:
            user: User identifier
            chunk_ids: List of chunk IDs to delete

        Returns:
            Number of chunks deleted
        """
        if not chunk_ids:
            return 0

        collection_name = get_qdrant_collection_name(user)

        try:
            deleted_count = await self.qdrant.delete_by_ids(
                collection_name=collection_name,
                point_ids=chunk_ids
            )

            logger.info(f"Purged {deleted_count} orphan chunks for user {user}")
            return deleted_count

        except Exception as e:
            logger.error(f"Failed to purge chunks: {e}", exc_info=True)
            return 0

    async def find_duplicate_pairs(
        self,
        user: str,
        similarity_threshold: float = 0.9,
        max_chunks_scanned: int = 2000,
        max_pairs: int = 100
    ) -> Dict[str, Any]:
        """
        Tenant-scoped similarity scan for near-duplicate wiki pages.

        Scrolls the tenant's own Qdrant collection (never another tenant's),
        then queries each chunk's vector against the same collection. Chunk
        pairs from DIFFERENT pages scoring above the threshold are grouped
        per page pair with the best score and the number of matching chunk
        pairs. Read-only: nothing is modified.

        Args:
            user: Tenant user identifier
            similarity_threshold: Minimum cosine similarity (default 0.9)
            max_chunks_scanned: Safety cap on chunks used as probes
            max_pairs: Maximum page pairs returned (highest score first)

        Returns:
            {
                "chunks_scanned": int,
                "duplicate_groups": [
                    {
                        "pages": [{page_id, path, title}, {page_id, path, title}],
                        "max_similarity": float,
                        "matching_chunk_pairs": int
                    }, ...
                ]
            }
        """
        collection_name = get_qdrant_collection_name(user)

        exists = await self.qdrant.collection_exists(collection_name)
        if not exists:
            return {"chunks_scanned": 0, "duplicate_groups": []}

        points = await self.qdrant.scroll_all_points(
            collection_name=collection_name,
            batch_size=100,
            with_payload=True,
            with_vectors=True
        )

        # Only wiki chunks participate (documents have their own dedup story)
        wiki_points = [
            p for p in points
            if p.get("vector") is not None
            and (p.get("payload") or {}).get("doc_type", "wiki") == "wiki"
            and (p.get("payload") or {}).get("page_id")
        ][:max_chunks_scanned]

        page_meta: Dict[int, Dict[str, Any]] = {}
        pair_stats: Dict[tuple, Dict[str, Any]] = {}
        seen_chunk_pairs = set()

        for point in wiki_points:
            payload = point.get("payload") or {}
            page_id = payload.get("page_id")
            page_meta.setdefault(page_id, {
                "page_id": page_id,
                "path": payload.get("page_path", ""),
                "title": payload.get("page_title", "")
            })

            hits = await self.qdrant.search_vectors(
                collection_name=collection_name,
                query_vector=point["vector"],
                limit=10,
                score_threshold=similarity_threshold
            )

            for hit in hits:
                hit_payload = hit.get("payload") or {}
                hit_page_id = hit_payload.get("page_id")
                if not hit_page_id or hit_page_id == page_id:
                    continue
                if hit_payload.get("doc_type", "wiki") != "wiki":
                    continue

                # Deduplicate the A->B / B->A chunk pair directions
                chunk_pair = tuple(sorted((point["id"], hit["id"])))
                if chunk_pair in seen_chunk_pairs:
                    continue
                seen_chunk_pairs.add(chunk_pair)

                page_meta.setdefault(hit_page_id, {
                    "page_id": hit_page_id,
                    "path": hit_payload.get("page_path", ""),
                    "title": hit_payload.get("page_title", "")
                })

                page_pair = tuple(sorted((page_id, hit_page_id)))
                stats = pair_stats.setdefault(page_pair, {
                    "max_similarity": 0.0,
                    "matching_chunk_pairs": 0
                })
                stats["max_similarity"] = max(stats["max_similarity"], hit["score"])
                stats["matching_chunk_pairs"] += 1

        groups = [
            {
                "pages": [page_meta[a], page_meta[b]],
                "max_similarity": stats["max_similarity"],
                "matching_chunk_pairs": stats["matching_chunk_pairs"]
            }
            for (a, b), stats in pair_stats.items()
        ]
        groups.sort(key=lambda g: g["max_similarity"], reverse=True)

        logger.info(
            f"Duplicate scan for {user}: {len(wiki_points)} chunks scanned, "
            f"{len(groups)} page pairs above {similarity_threshold}"
        )

        return {
            "chunks_scanned": len(wiki_points),
            "duplicate_groups": groups[:max_pairs]
        }

    def find_chunks_without_graph_nodes(
        self,
        chunk_references: List[Dict[str, Any]],
        graph_references: List[Dict[str, Any]]
    ) -> List[str]:
        """
        Find vector chunks that have no corresponding graph Document node.

        Used for bidirectional orphan detection - vectors without graph representation.

        Args:
            chunk_references: List from get_all_chunk_references()
            graph_references: List from GraphService.get_all_document_references()

        Returns:
            List of orphan chunk IDs
        """
        # Build sets of IDs that have graph nodes
        graph_page_ids = {
            ref.get("page_id") for ref in graph_references
            if ref.get("doc_type") == "wiki" and ref.get("page_id")
        }
        graph_doc_ids = {
            ref.get("document_id") for ref in graph_references
            if ref.get("doc_type") != "wiki" and ref.get("document_id")
        }

        # Find chunks with no graph node
        orphan_ids = []
        for chunk in chunk_references:
            doc_type = chunk.get("doc_type", "wiki")

            if doc_type == "wiki":
                page_id = chunk.get("page_id")
                if page_id and page_id not in graph_page_ids:
                    orphan_ids.append(chunk["chunk_id"])
            else:
                document_id = chunk.get("document_id")
                if document_id and document_id not in graph_doc_ids:
                    orphan_ids.append(chunk["chunk_id"])

        logger.info(f"Found {len(orphan_ids)} vector chunks without graph nodes")
        return orphan_ids