"""
BM25 keyword-based search for hybrid retrieval.
Provides exact keyword matching to complement semantic search.
"""
from typing import List, Dict
from rank_bm25 import BM25Okapi
import numpy as np


class BM25Search:
    """BM25 keyword-based search engine."""
    
    def __init__(self):
        """Initialize BM25 search."""
        self.corpus = []
        self.tokenized_corpus = []
        self.bm25 = None
        self.chunks = []
        print("[BM25] Initialized")
    
    def index_documents(self, chunks: List[Dict]):
        """
        Index documents for BM25 search.
        
        Args:
            chunks: List of chunk dictionaries with 'content' key
        """
        self.chunks = chunks
        self.corpus = [chunk['content'] for chunk in chunks]
        
        # Simple tokenization (lowercase + split)
        self.tokenized_corpus = [
            doc.lower().split() for doc in self.corpus
        ]
        
        # Build BM25 index
        if self.tokenized_corpus:
            self.bm25 = BM25Okapi(self.tokenized_corpus)
            print(f"[BM25] Indexed {len(self.chunks)} documents")
    
    def search(
        self,
        query: str,
        top_k: int = 30
    ) -> List[Dict]:
        """
        Search documents using BM25.
        
        Args:
            query: Search query
            top_k: Number of results to return
            
        Returns:
            List of chunks sorted by BM25 score
        """
        if not self.bm25:
            return []
        
        # Tokenize query
        tokenized_query = query.lower().split()
        
        # Get BM25 scores
        scores = self.bm25.get_scores(tokenized_query)
        
        # Get top_k indices
        top_indices = np.argsort(scores)[::-1][:top_k]
        
        # Return chunks with scores
        results = []
        for idx in top_indices:
            if scores[idx] > 0:  # Only return non-zero scores
                chunk = self.chunks[idx].copy()
                chunk['bm25_score'] = float(scores[idx])
                results.append(chunk)
        
        return results
    
    def get_score(self, query: str, document: str) -> float:
        """
        Get BM25 score for a single query-document pair.
        
        Args:
            query: Search query
            document: Document text
            
        Returns:
            BM25 relevance score
        """
        tokenized_query = query.lower().split()
        tokenized_doc = document.lower().split()
        
        # Create temporary BM25 for single document
        temp_bm25 = BM25Okapi([tokenized_doc])
        score = temp_bm25.get_scores(tokenized_query)[0]
        
        return float(score)


class HybridSearch:
    """Combines semantic and keyword search."""
    
    def __init__(
        self,
        semantic_weight: float = 0.7,
        keyword_weight: float = 0.3
    ):
        """
        Initialize hybrid search.
        
        Args:
            semantic_weight: Weight for semantic search (0-1)
            keyword_weight: Weight for keyword search (0-1)
        """
        self.semantic_weight = semantic_weight
        self.keyword_weight = keyword_weight
        self.bm25 = BM25Search()
        print(f"[Hybrid Search] Initialized (semantic: {semantic_weight}, keyword: {keyword_weight})")
    
    def combine_results(
        self,
        semantic_results: List[Dict],
        keyword_results: List[Dict],
        top_k: int = 10
    ) -> List[Dict]:
        """
        Combine and rerank results from semantic and keyword search.
        
        Args:
            semantic_results: Results from semantic search (with 'distance' scores)
            keyword_results: Results from BM25 search (with 'bm25_score')
            top_k: Number of final results
            
        Returns:
            Combined and reranked results
        """
        # Normalize scores to 0-1 range
        def normalize_scores(results, score_key):
            if not results:
                return results
            
            scores = [r.get(score_key, 0) for r in results]
            min_score = min(scores)
            max_score = max(scores)
            
            if max_score == min_score:
                return results
            
            for r in results:
                r[f'{score_key}_normalized'] = (
                    (r.get(score_key, 0) - min_score) / (max_score - min_score)
                )
            return results
        
        # For semantic search, lower distance = higher relevance
        # Need to invert: score = 1 - normalized_distance
        for r in semantic_results:
            if 'distance' in r:
                r['semantic_score'] = r['distance']  # Will normalize below
        
        semantic_results = normalize_scores(semantic_results, 'semantic_score')
        keyword_results = normalize_scores(keyword_results, 'bm25_score')
        
        # Invert semantic scores (lower distance = better)
        for r in semantic_results:
            if 'semantic_score_normalized' in r:
                r['semantic_score_normalized'] = 1 - r['semantic_score_normalized']
        
        # Merge results by chunk ID or content
        merged = {}
        
        for chunk in semantic_results:
            chunk_id = chunk.get('metadata', {}).get('document_id', '') + '_' + str(chunk.get('metadata', {}).get('chunk_index', ''))
            merged[chunk_id] = chunk.copy()
            merged[chunk_id]['hybrid_score'] = (
                self.semantic_weight * chunk.get('semantic_score_normalized', 0)
            )
        
        for chunk in keyword_results:
            chunk_id = chunk.get('metadata', {}).get('document_id', '') + '_' + str(chunk.get('metadata', {}).get('chunk_index', ''))
            if chunk_id in merged:
                merged[chunk_id]['hybrid_score'] += (
                    self.keyword_weight * chunk.get('bm25_score_normalized', 0)
                )
            else:
                merged[chunk_id] = chunk.copy()
                merged[chunk_id]['hybrid_score'] = (
                    self.keyword_weight * chunk.get('bm25_score_normalized', 0)
                )
        
        # Sort by hybrid score
        results = sorted(
            merged.values(),
            key=lambda x: x.get('hybrid_score', 0),
            reverse=True
        )
        
        return results[:top_k]


# Global hybrid search instance
hybrid_search = HybridSearch()