src/rag/vector_store.py

"""
Qdrant Vector Store
Manages vector embeddings in Qdrant database
"""

import os
import sys
import threading
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from typing import List, Dict

from src.config.paths import QDRANT_DIR


class QdrantVectorStore:
    """
    Qdrant Vector Store Manager
    
    Handles storage and retrieval of document embeddings
    
    Usage:
        store = QdrantVectorStore()
        store.add_document_chunk("chunk_id", embedding, metadata)
        results = store.search_documents(query_embedding)
    """
    
    def __init__(self, host: str = "localhost", port: int = 6333):
        """
        Initialize Qdrant client
        
        Args:
            host: Qdrant host (default: localhost)
            port: Qdrant port (default: 6333)
        """
        self.host = host
        self.port = port
        self.embedding_dim = 1536  # OpenAI text-embedding-3-small
        self.using_docker = False
        # Serializes access to the client. The embedded (on-disk) Qdrant client is
        # not safe under concurrent threads, so search/upsert take this lock. Paired
        # with the process-wide singleton (get_shared_vector_store) this makes the
        # multi-user Space safe: one client, one-at-a-time access.
        self._lock = threading.RLock()

        # QDRANT_LOCAL_PATH env still overrides; default is the project-root
        # anchored QDRANT_DIR (build scratch). Empty string forces in-memory.
        local_path_env = os.environ.get("QDRANT_LOCAL_PATH")
        local_path = local_path_env if local_path_env is not None else str(QDRANT_DIR)

        # Build mode: write directly to the on-disk store and skip the Docker
        # probe (used by scripts/build_index.py). Avoids accidentally connecting
        # to a stray empty Docker on :6333 during a build.
        if os.environ.get("QDRANT_FORCE_LOCAL", "0") == "1":
            self.client = QdrantClient(path=local_path)
            print(f"📦 QDRANT_FORCE_LOCAL: using on-disk store at '{local_path}' (build mode)")
            self.initialize_collections()
            return

        # Served mode: try Docker Qdrant first (the deploy restores the committed
        # snapshot into it and both processes read from it).
        try:
            self.client = QdrantClient(host=host, port=port, timeout=5)
            # Test connection by getting collections
            _ = self.client.get_collections()
            self.using_docker = True
            print(f"✅ Connected to Qdrant Docker at {host}:{port}")
            self.initialize_collections()
            return
        except Exception as e:
            docker_error = e

        # Fail fast in the hosted deploy: a served instance MUST read the Docker
        # Qdrant restored from the committed snapshot. Silently falling back to an
        # empty in-memory (or fresh on-disk) store would answer nothing or trigger
        # a full re-embed. Opt in with QDRANT_REQUIRE_DOCKER=1 in the deploy env.
        if os.environ.get("QDRANT_REQUIRE_DOCKER", "0") == "1":
            raise RuntimeError(
                f"QDRANT_REQUIRE_DOCKER=1 but no Qdrant reachable at {host}:{port}. "
                f"Start Docker Qdrant and restore the committed snapshot "
                f"(data/qdrant_snapshots/, see RUNBOOK). Original error: {docker_error}"
            )

        # Legacy / dev fallback. Persistent on-disk store so embeddings survive
        # restarts (fixes the cold-start re-embed ISA reported); in-memory only
        # if the path is unavailable or explicitly disabled (QDRANT_LOCAL_PATH="").
        print(f"⚠️  Qdrant Docker not available at {host}:{port}")
        print(f"   Error: {str(docker_error)}")
        print("   💡 For production, start Docker: docker run -d -p 6333:6333 qdrant/qdrant")
        if local_path:
            try:
                self.client = QdrantClient(path=local_path)
                print(f"   📦 Using persistent on-disk store at '{local_path}' (survives restarts)")
            except Exception as e2:
                print(f"   ⚠️  On-disk store unavailable ({e2}); using in-memory (data lost on restart)")
                self.client = QdrantClient(":memory:")
        else:
            print("   📦 Using in-memory mode (data lost on restart)")
            self.client = QdrantClient(":memory:")
        self.initialize_collections()
    
    def initialize_collections(self):
        """Create collections for APIs and PDF documents"""
        # Collection for API descriptions (for future use)
        if not self.client.collection_exists("api_catalog"):
            try:
                self.client.create_collection(
                    collection_name="api_catalog",
                    vectors_config=VectorParams(
                        size=self.embedding_dim,
                        distance=Distance.COSINE
                    )
                )
            except Exception:
                pass  # Collection might already exist
        
        # Collection for PDF document chunks (CDMS)
        if not self.client.collection_exists("cdms_documents"):
            try:
                self.client.create_collection(
                    collection_name="cdms_documents",
                    vectors_config=VectorParams(
                        size=self.embedding_dim,
                        distance=Distance.COSINE
                    )
                )
            except Exception:
                pass  # Collection might already exist
    
    def add_document_chunk(
        self, 
        chunk_id: str, 
        embedding: List[float], 
        metadata: Dict
    ):
        """
        Add PDF chunk to vector store
        
        Args:
            chunk_id: Unique identifier for the chunk (string, will be converted to int)
            embedding: Vector embedding (1536 dimensions)
            metadata: Dict with chunk metadata
        """
        try:
            # Convert string ID to integer hash for Qdrant (Qdrant requires int or UUID)
            import hashlib
            int_id = int(hashlib.md5(chunk_id.encode()).hexdigest()[:15], 16)  # Use first 15 hex chars as int

            with self._lock:
                self.client.upsert(
                    collection_name="cdms_documents",
                    points=[PointStruct(
                        id=int_id,  # Qdrant requires integer or UUID
                        vector=embedding,
                        payload=metadata
                    )]
                )
        except Exception as e:
            print(f"⚠️  Warning: Could not add chunk to Qdrant: {e}")
    
    def search_documents(
        self,
        query_embedding: List[float],
        limit: int = 10,
        score_threshold: float = 0.4,  # Balanced: 0.3 let weak/wrong-product
        # matches through (they were being answered from); 0.5 was too strict.
        # Tune with real queries. Better to abstain than cite the wrong label.
        query_filter=None,
    ) -> List[Dict]:
        """
        Search for similar document chunks

        Args:
            query_embedding: Query vector embedding
            limit: Maximum number of results
            score_threshold: Minimum similarity score (0-1)
            query_filter: Optional Qdrant Filter to scope the search (e.g. by
                product) so retrieval is constrained at the vector level rather
                than post-filtered from a global, dominant-product-heavy top-k.

        Returns:
            List of search results with metadata and scores
        """
        try:
            with self._lock:
                results = self.client.search(
                    collection_name="cdms_documents",
                    query_vector=query_embedding,
                    query_filter=query_filter,
                    limit=limit,
                    score_threshold=score_threshold,
                    with_payload=True
                )
            
            # Format results
            formatted_results = []
            for result in results:
                payload = result.payload
                formatted_results.append({
                    "id": result.id,
                    "score": result.score,
                    "content": payload.get("content", ""),
                    "source_file": payload.get("source_file", "Unknown"),
                    "page_number": payload.get("page_number", 0),
                    "document_id": payload.get("document_id", ""),
                    "pdf_url": payload.get("pdf_url", ""),  # PHASE 1 FIX: Extract PDF URL from payload
                    "url_hash": payload.get("url_hash", ""),  # PHASE 1 FIX: Extract URL hash from payload
                    "metadata": payload
                })
            
            return formatted_results
        
        except Exception as e:
            print(f"⚠️  Warning: Could not search Qdrant: {e}")
            return []
    
    def get_collection_info(self) -> Dict:
        """Get information about collections"""
        try:
            info = {
                "using_docker": self.using_docker,
                "host": self.host if self.using_docker else "in-memory",
                "port": self.port if self.using_docker else None
            }
            if self.client.collection_exists("cdms_documents"):
                collection_info = self.client.get_collection("cdms_documents")
                info["cdms_documents"] = {
                    "points_count": collection_info.points_count,
                    # vectors_count was removed in qdrant-client >=1.16; fall back
                    # to points_count so this stays version-tolerant.
                    "vectors_count": getattr(collection_info, "vectors_count", None)
                    or collection_info.points_count
                }
            return info
        except Exception as e:
            return {"error": str(e)}


# --- Process-wide singleton --------------------------------------------------
# The embedded (on-disk) Qdrant permits only ONE client per path per process: a
# second QdrantClient(path=...) raises "already accessed by another instance".
# Every request previously built its own store, so two concurrent users on the
# Space would collide -> the loser fell back to an empty in-memory store and
# abstained. Share ONE store across all requests (double-checked lock for a safe
# first init); per-op locking inside the store serialises concurrent access.
_SHARED_VS = None
_SHARED_VS_LOCK = threading.Lock()


def get_shared_vector_store() -> "QdrantVectorStore":
    """Return the process-wide QdrantVectorStore, creating it once on first use."""
    global _SHARED_VS
    if _SHARED_VS is None:
        with _SHARED_VS_LOCK:
            if _SHARED_VS is None:
                _SHARED_VS = QdrantVectorStore()
    return _SHARED_VS


# Test function
if __name__ == "__main__":
    print("Testing Qdrant Vector Store...")
    print("-" * 70)
    
    try:
        store = QdrantVectorStore()
        
        print("✅ Qdrant connection successful!")
        
        info = store.get_collection_info()
        if info:
            print("\n📊 Collection Info:")
            for collection, data in info.items():
                print(f"   {collection}: {data.get('points_count', 0)} points")
        else:
            print("\n📊 Collections initialized (empty)")
        
        print("\n💡 To use:")
        print("   1. Make sure Qdrant is running:")
        print("      docker run -d -p 6333:6333 qdrant/qdrant")
        print("   2. Or it will use in-memory mode automatically")
        
    except Exception as e:
        print(f"❌ Error: {e}")