app/services/vector_store.py

import hashlib
import json
from pathlib import Path
from typing import Dict, List, Tuple

import faiss
import numpy as np
import pickle
from rank_bm25 import BM25Okapi

from app.services.chunker import chunk_documents
from app.services.document_loader import load_documents
from app.services.embeddings import EmbeddingService
import re
import asyncio

def _analyze(text: str) -> List[str]:
    """Professional analyzer for RAG: lowercases, removes noise, and handles hyphens/punctuation."""
    # 1. Lowercase
    text = text.lower()
    
    # 2. Handle compound words: index "pro-rata" or "half/day" as joined "prorata", "halfday"
    # Find all words containing hyphens or slashes
    compounds = re.findall(r'\b\w+(?:[-\/]\w+)+\b', text)
    for word in compounds:
        joined = word.replace('-', '').replace('/', '')
        text += f" {joined}"
        
    # 3. Final tokenization: split by non-word characters to remove punctuation but keep words
    tokens = re.findall(r'\b\w+\b', text)
    
    # 4. Filter short noise tokens (optional, but keep for precision)
    return [t for t in tokens if len(t) > 1]


class FaissVectorStore:
    def __init__(
        self,
        embedding_service: EmbeddingService,
        docs_dir: Path,
        index_dir: Path,
        chunk_size_tokens: int,
        chunk_overlap_tokens: int,
    ) -> None:
        self.embedding_service = embedding_service
        self.docs_dir = docs_dir
        self.index_dir = index_dir
        self.chunk_size_tokens = chunk_size_tokens
        self.chunk_overlap_tokens = chunk_overlap_tokens

        self.faiss_index_path = index_dir / "faiss.index"
        self.bm25_index_path = index_dir / "bm25.pkl"
        self.metadata_path = index_dir / "metadata.json"
        self.state_path = index_dir / "state.json"

        self.index = None
        self.bm25 = None
        self.metadata: List[Dict[str, str]] = []
        self.docs_loaded = False
        self.last_retrieved: List[Dict[str, str]] = []

    def _compute_docs_fingerprint(self) -> str:
        hasher = hashlib.sha256()
        # Include analyzer version and chunk settings in fingerprint so changing them triggers re-index
        hasher.update("v4_hybrid_normalizer".encode("utf-8"))  # bump this to force re-index
        hasher.update(str(self.chunk_size_tokens).encode("utf-8"))
        hasher.update(str(self.chunk_overlap_tokens).encode("utf-8"))
        
        if not self.docs_dir.exists():
            return "no_docs"
        for path in sorted(self.docs_dir.rglob("*")):
            if path.is_file() and path.suffix.lower() in {".txt", ".pdf"}:
                stat = path.stat()
                hasher.update(str(path).encode("utf-8"))
                hasher.update(str(stat.st_mtime_ns).encode("utf-8"))
                hasher.update(str(stat.st_size).encode("utf-8"))
        return hasher.hexdigest()

    def _can_use_cached_index(self, fingerprint: str) -> bool:
        if not (self.faiss_index_path.exists() and self.metadata_path.exists() and self.state_path.exists()):
            return False
        try:
            state = json.loads(self.state_path.read_text(encoding="utf-8"))
            return state.get("docs_fingerprint") == fingerprint
        except Exception:
            return False

    def build_or_load(self) -> None:
        self.index_dir.mkdir(parents=True, exist_ok=True)
        fingerprint = self._compute_docs_fingerprint()

        if self._can_use_cached_index(fingerprint):
            self.index = faiss.read_index(str(self.faiss_index_path))
            with open(self.bm25_index_path, "rb") as f:
                self.bm25 = pickle.load(f)
            self.metadata = json.loads(self.metadata_path.read_text(encoding="utf-8"))
            self.docs_loaded = len(self.metadata) > 0
            return

        docs = load_documents(self.docs_dir)
        chunks = chunk_documents(docs, self.chunk_size_tokens, self.chunk_overlap_tokens)
        if not chunks:
            self.index = None
            self.metadata = []
            self.docs_loaded = False
            return

        vectors = self.embedding_service.encode([c["text"] for c in chunks])
        dim = vectors.shape[1]
        index = faiss.IndexFlatIP(dim)
        index.add(vectors)

        tokenized_corpus = [_analyze(c["text"]) for c in chunks]
        bm25 = BM25Okapi(tokenized_corpus)

        self.index = index
        self.bm25 = bm25
        self.metadata = chunks
        self.docs_loaded = True

        faiss.write_index(index, str(self.faiss_index_path))
        with open(self.bm25_index_path, "wb") as f:
            pickle.dump(bm25, f)
        self.metadata_path.write_text(json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8")
        self.state_path.write_text(
            json.dumps({"docs_fingerprint": fingerprint, "chunk_count": len(chunks)}, indent=2),
            encoding="utf-8",
        )

    def search(self, query: str, top_k: int = 4) -> List[Dict[str, str]]:
        """Backwards compatibility for single query search."""
        return self.multi_search([query], top_k=top_k)

    def multi_search(self, queries: List[str], top_k: int = 4) -> List[Dict[str, str]]:
        """
        Professional batched search:
        1. Encodes all queries in a single batch (fast).
        2. Searches FAISS for all vectors at once.
        3. Runs BM25 for all queries.
        4. Combines everything with a unified RRF pass.
        """
        if self.index is None or not self.metadata or self.bm25 is None or not queries:
            self.last_retrieved = []
            return []

        # 1. Batched Embedding — use encode() which handles any list size efficiently
        query_vectors = self.embedding_service.encode(list(queries))
        # Ensure shape is always 2D: (N, dim)
        if query_vectors.ndim == 1:
            query_vectors = query_vectors.reshape(1, -1)
        
        # 2. Batched FAISS Search
        # faiss_indices shape: (len(queries), top_k*2)
        faiss_scores, faiss_indices = self.index.search(
            np.asarray(query_vectors, dtype=np.float32), 
            top_k * 2  # Reasonable pool for RRF merging
        )

        # 3. Batched BM25 Search
        # Combine all tokenized queries
        k = 60
        rrf_scores = {}

        for q_idx, query in enumerate(queries):
            # FAISS results for this query
            for rank, idx in enumerate(faiss_indices[q_idx]):
                if idx < 0 or idx >= len(self.metadata):
                    continue
                # Add to RRF score
                rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (1.0 / (k + rank + 1))
            
            # BM25 results for this query
            tokenized_query = _analyze(query)
            bm25_scores = self.bm25.get_scores(tokenized_query)
            bm25_top_indices = np.argsort(bm25_scores)[::-1][:top_k * 2]
            
            for rank, idx in enumerate(bm25_top_indices):
                if idx < 0 or idx >= len(self.metadata) or bm25_scores[idx] <= 0:
                    continue
                rrf_scores[idx] = rrf_scores.get(idx, 0.0) + (1.0 / (k + rank + 1))

        # 4. Final Ranking
        if not rrf_scores:
            self.last_retrieved = []
            return []

        sorted_indices = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)[:top_k]

        results: List[Dict[str, str]] = []
        for idx in sorted_indices:
            chunk = self.metadata[idx]
            results.append(
                {
                    "id": chunk["id"],
                    "source": chunk["source"],
                    "text": chunk["text"],
                    "score": float(rrf_scores[idx]),
                }
            )
        
        self.last_retrieved = results
        return results

    def health(self) -> Dict[str, object]:
        return {
            "docs_loaded": self.docs_loaded,
            "index_ready": self.index is not None,
            "chunk_count": len(self.metadata),
            "index_path": str(self.faiss_index_path),
        }