src/retriever.py

"""
Hybrid Retriever
-----------------
Combines semantic search (FAISS) and keyword search (BM25) using 
Reciprocal Rank Fusion (RRF) for robust retrieval.

Why hybrid?
- Semantic search catches meaning ("investigative reporting" → "accountability journalism")
- BM25 catches exact terms ("Pulitzer Center", "Africa", specific deadlines)
- RRF merges both rankings without needing tuned weights

Additional features:
- Metadata-aware filtering (deadline, region, type)
- Query classification to route questions optimally
- Source deduplication across retrievers
"""

import re
from datetime import datetime, timedelta
from typing import Optional

from langchain_core.documents import Document
from langchain_community.vectorstores import FAISS
from rank_bm25 import BM25Okapi


# ---------------------------------------------------------------------------
# QUERY CLASSIFIER
# ---------------------------------------------------------------------------

def classify_query(query: str) -> dict:
    """
    Classify the user query to determine optimal retrieval strategy.
    
    Returns a dict with:
      - intent: general | deadline_search | region_search | topic_search | newsletter | about
      - filters: any metadata filters to apply
      - boost_keywords: terms to emphasize in BM25
    """
    q_lower = query.lower()
    result = {"intent": "general", "filters": {}, "boost_keywords": []}
    
    # --- Deadline-based queries ---
    deadline_patterns = [
        r"deadline.*(next|coming|within)\s+(\d+)\s*(day|week|month)",
        r"(next|coming|within)\s+(\d+)\s*(day|week|month)",
        r"due\s+soon", r"closing\s+soon", r"expiring",
        r"deadline", r"when.*due",
    ]
    for pattern in deadline_patterns:
        if re.search(pattern, q_lower):
            result["intent"] = "deadline_search"
            # Try to extract the time window
            match = re.search(r"(\d+)\s*(day|week|month)", q_lower)
            if match:
                num = int(match.group(1))
                unit = match.group(2)
                if unit.startswith("week"):
                    num *= 7
                elif unit.startswith("month"):
                    num *= 30
                result["filters"]["deadline_days"] = num
            else:
                result["filters"]["deadline_days"] = 30  # default
            break
    
    # --- Region-based queries ---
    region_map = {
        "africa": ["Africa", "Sub-Saharan Africa", "East Africa", "North Africa"],
        "asia": ["Asia-Pacific", "South Asia", "Southeast Asia", "East Asia"],
        "latin america": ["Latin America", "Central America", "South America"],
        "europe": ["Europe", "EU", "Eastern Europe"],
        "middle east": ["Middle East", "MENA", "North Africa"],
        "mena": ["Middle East", "MENA", "North Africa"],
        "south asia": ["South Asia", "India", "Pakistan", "Bangladesh"],
        "india": ["South Asia", "India"],
    }
    for keyword, regions in region_map.items():
        if keyword in q_lower:
            result["intent"] = "region_search"
            result["filters"]["regions"] = regions
            result["boost_keywords"].extend(regions)
            break
    
    # --- Type-based queries ---
    type_map = {
        "fellowship": "fellowship",
        "grant": "grant",
        "training": "training",
        "award": "award",
        "workshop": "training",
    }
    for keyword, opp_type in type_map.items():
        if keyword in q_lower:
            result["filters"]["opp_type"] = opp_type
            result["boost_keywords"].append(opp_type)
            break
    
    # --- Newsletter queries ---
    if any(w in q_lower for w in ["newsletter", "subscribe", "email list", "mailing list"]):
        result["intent"] = "newsletter"
        result["boost_keywords"].extend(["newsletter", "subscribe"])
    
    # --- About IJNet ---
    if any(w in q_lower for w in ["what is ijnet", "about ijnet", "what does ijnet"]):
        result["intent"] = "about"
    
    # --- Topic detection ---
    topic_keywords = {
        "ai": ["AI tools", "artificial intelligence", "AI"],
        "data journalism": ["data journalism", "data analysis"],
        "investigative": ["investigative journalism", "investigation"],
        "climate": ["climate change", "environment"],
        "fact-check": ["fact-checking", "verification"],
        "digital security": ["digital security", "journalist safety"],
        "product": ["product design", "newsroom innovation", "UX"],
        "design": ["product design", "UX", "design"],
        "solutions journalism": ["solutions journalism", "constructive"],
        "mobile journalism": ["mobile journalism", "MoJo"],
        "freelance": ["freelance", "independent reporting"],
        "press freedom": ["press freedom", "media freedom"],
    }
    for keyword, topics in topic_keywords.items():
        if keyword in q_lower:
            result["boost_keywords"].extend(topics)
    
    return result


# ---------------------------------------------------------------------------
# BM25 INDEX
# ---------------------------------------------------------------------------

class BM25Index:
    """Lightweight BM25 index over documents for keyword retrieval."""
    
    def __init__(self, documents: list[Document]):
        self.documents = documents
        # Tokenize for BM25
        self.tokenized = [
            self._tokenize(doc.page_content) for doc in documents
        ]
        self.bm25 = BM25Okapi(self.tokenized)
    
    @staticmethod
    def _tokenize(text: str) -> list[str]:
        """Simple whitespace + punctuation tokenizer with lowercasing."""
        text = text.lower()
        text = re.sub(r"[^\w\s]", " ", text)
        return text.split()
    
    def search(self, query: str, k: int = 10) -> list[tuple[Document, float]]:
        """Return top-k documents with BM25 scores."""
        tokenized_query = self._tokenize(query)
        scores = self.bm25.get_scores(tokenized_query)
        
        # Get top-k indices
        top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
        
        return [(self.documents[i], scores[i]) for i in top_indices if scores[i] > 0]


# ---------------------------------------------------------------------------
# HYBRID RETRIEVER
# ---------------------------------------------------------------------------

class HybridRetriever:
    """
    Combines FAISS (semantic) and BM25 (keyword) retrieval with
    Reciprocal Rank Fusion and metadata-aware post-filtering.
    """
    
    def __init__(
        self,
        vector_store: FAISS,
        documents: list[Document],
        semantic_k: int = 8,
        bm25_k: int = 8,
        final_k: int = 5,
        rrf_constant: int = 60,
    ):
        self.vector_store = vector_store
        self.bm25_index = BM25Index(documents)
        self.all_documents = documents
        self.semantic_k = semantic_k
        self.bm25_k = bm25_k
        self.final_k = final_k
        self.rrf_constant = rrf_constant  # standard RRF constant
    
    def _reciprocal_rank_fusion(
        self,
        semantic_results: list[Document],
        bm25_results: list[tuple[Document, float]],
    ) -> list[Document]:
        """
        Merge two ranked lists using RRF.
        
        RRF score = Σ 1 / (k + rank)  across all lists where the doc appears.
        This is robust to score magnitude differences between retrievers.
        """
        doc_scores: dict[str, float] = {}
        doc_map: dict[str, Document] = {}
        
        # Score semantic results
        for rank, doc in enumerate(semantic_results):
            doc_id = self._doc_key(doc)
            score = 1.0 / (self.rrf_constant + rank + 1)
            doc_scores[doc_id] = doc_scores.get(doc_id, 0) + score
            doc_map[doc_id] = doc
        
        # Score BM25 results
        for rank, (doc, _bm25_score) in enumerate(bm25_results):
            doc_id = self._doc_key(doc)
            score = 1.0 / (self.rrf_constant + rank + 1)
            doc_scores[doc_id] = doc_scores.get(doc_id, 0) + score
            doc_map[doc_id] = doc
        
        # Sort by fused score
        sorted_ids = sorted(doc_scores, key=doc_scores.get, reverse=True)
        return [doc_map[did] for did in sorted_ids]
    
    @staticmethod
    def _doc_key(doc: Document) -> str:
        """Generate a unique key for deduplication."""
        doc_id = doc.metadata.get("doc_id", "")
        chunk_idx = doc.metadata.get("chunk_index", 0)
        return f"{doc_id}_chunk{chunk_idx}"
    
    def _apply_metadata_filters(
        self,
        documents: list[Document],
        filters: dict,
    ) -> list[Document]:
        """
        Post-retrieval metadata filtering.
        
        Moves matching documents to the top rather than removing non-matches,
        so we still return results even if filters don't match perfectly.
        """
        if not filters:
            return documents
        
        matching = []
        non_matching = []
        
        for doc in documents:
            match = True
            
            # Deadline filter
            if "deadline_days" in filters:
                deadline_str = doc.metadata.get("deadline", "")
                if deadline_str:
                    try:
                        deadline = datetime.strptime(deadline_str, "%Y-%m-%d")
                        cutoff = datetime.now() + timedelta(days=filters["deadline_days"])
                        if deadline > cutoff or deadline < datetime.now():
                            match = False
                    except ValueError:
                        pass  # Can't parse date, don't filter
            
            # Region filter
            if "regions" in filters:
                doc_regions = doc.metadata.get("regions", "").lower()
                if doc_regions and not any(r.lower() in doc_regions for r in filters["regions"]):
                    match = False
            
            # Type filter
            if "opp_type" in filters:
                doc_type = doc.metadata.get("opp_type", "").lower()
                if doc_type and doc_type != filters["opp_type"].lower():
                    match = False
            
            if match:
                matching.append(doc)
            else:
                non_matching.append(doc)
        
        # Return matching first, then non-matching as fallback
        return matching + non_matching
    
    def retrieve(self, query: str) -> list[Document]:
        """
        Full hybrid retrieval pipeline:
        1. Classify query
        2. Run semantic + BM25 search (with boosted keywords)
        3. Fuse with RRF
        4. Apply metadata filters
        5. Return top-k unique results
        """
        # Step 1: Classify
        classification = classify_query(query)
        
        # Step 2a: Augment query with boost keywords for better BM25
        augmented_query = query
        if classification["boost_keywords"]:
            augmented_query = query + " " + " ".join(classification["boost_keywords"])
        
        # Step 2b: Semantic search
        semantic_results = self.vector_store.similarity_search(
            query, k=self.semantic_k
        )
        
        # Step 2c: BM25 search (with augmented query)
        bm25_results = self.bm25_index.search(augmented_query, k=self.bm25_k)
        
        # Step 3: Reciprocal Rank Fusion
        fused = self._reciprocal_rank_fusion(semantic_results, bm25_results)
        
        # Step 4: Metadata-aware reranking
        filtered = self._apply_metadata_filters(fused, classification["filters"])
        
        # Step 5: Return top-k
        return filtered[:self.final_k]
    
    def retrieve_with_debug(self, query: str) -> dict:
        """
        Same as retrieve() but returns debug info for development.
        """
        classification = classify_query(query)
        
        augmented_query = query
        if classification["boost_keywords"]:
            augmented_query = query + " " + " ".join(classification["boost_keywords"])
        
        semantic_results = self.vector_store.similarity_search_with_score(
            query, k=self.semantic_k
        )
        bm25_results = self.bm25_index.search(augmented_query, k=self.bm25_k)
        
        # Convert semantic results for fusion (drop scores for RRF)
        semantic_docs = [doc for doc, score in semantic_results]
        fused = self._reciprocal_rank_fusion(semantic_docs, bm25_results)
        filtered = self._apply_metadata_filters(fused, classification["filters"])
        
        return {
            "query": query,
            "classification": classification,
            "semantic_results": [(doc.metadata.get("title", ""), score) for doc, score in semantic_results[:5]],
            "bm25_results": [(doc.metadata.get("title", ""), score) for doc, score in bm25_results[:5]],
            "final_results": filtered[:self.final_k],
        }