"""
MLL Intake Embeddings Module
Vector embeddings and semantic similarity search
"""

import pickle
from pathlib import Path
from typing import List, Dict, Any, Optional
from loguru import logger

import sys
sys.path.insert(0, str(Path(__file__).parent.parent))
from config import settings

# Try to import numpy
try:
    import numpy as np
    NUMPY_AVAILABLE = True
except ImportError:
    NUMPY_AVAILABLE = False
    np = None
    logger.warning("NumPy not available")


class EmbeddingsManager:
    """Manage embeddings and vector search"""
    
    def __init__(self):
        self.model = None
        self.index = None
        self.documents: List[Dict[str, Any]] = []
        self.doc_ids: List[str] = []
        self._initialized = False
    
    def initialize(self, documents: Optional[List[Dict[str, Any]]] = None) -> bool:
        """Initialize the embedding model and optionally build index"""
        try:
            # Try to load sentence transformers
            try:
                from sentence_transformers import SentenceTransformer
                logger.info(f"Loading embedding model: {settings.vector.embedding_model}")
                self.model = SentenceTransformer(settings.vector.embedding_model)
            except ImportError:
                logger.warning("sentence-transformers not available, using simple text matching")
                self.model = None
            
            if documents and self.model:
                self.build_index(documents)
            elif documents:
                # Store documents for simple matching
                self.documents = documents
                self.doc_ids = [doc['id'] for doc in documents]
            
            self._initialized = True
            logger.info("Embeddings manager initialized")
            return True
            
        except Exception as e:
            logger.error(f"Error initializing embeddings: {e}")
            self._initialized = True  # Mark as initialized anyway
            return True
    
    def build_index(self, documents: List[Dict[str, Any]]) -> bool:
        """Build FAISS index from documents"""
        try:
            import faiss
            
            if not self.model:
                logger.warning("No model available for building index")
                self.documents = documents
                self.doc_ids = [doc['id'] for doc in documents]
                return True
            
            self.documents = documents
            self.doc_ids = [doc['id'] for doc in documents]
            texts = [doc['text'] for doc in documents]
            
            logger.info(f"Generating embeddings for {len(texts)} documents...")
            embeddings = self.model.encode(texts, show_progress_bar=False)
            
            # Normalize for cosine similarity
            embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
            
            # Build FAISS index
            dimension = embeddings.shape[1]
            self.index = faiss.IndexFlatIP(dimension)
            self.index.add(embeddings.astype(np.float32))
            
            logger.info(f"Built FAISS index with {self.index.ntotal} vectors")
            return True
            
        except Exception as e:
            logger.error(f"Error building index: {e}")
            return False
    
    def search(
        self, 
        query: str, 
        top_k: Optional[int] = None,
        threshold: Optional[float] = None
    ) -> List[Dict[str, Any]]:
        """Search for similar documents"""
        top_k = top_k or settings.vector.top_k
        threshold = threshold or settings.vector.similarity_threshold
        
        if not self._initialized:
            return []
        
        # If we have a proper index, use it
        if self.index is not None and self.model is not None:
            try:
                query_embedding = self.model.encode([query])
                query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
                
                scores, indices = self.index.search(query_embedding.astype(np.float32), min(top_k, self.index.ntotal))
                
                results = []
                for score, idx in zip(scores[0], indices[0]):
                    if idx < 0 or idx >= len(self.documents):
                        continue
                    if score < threshold:
                        continue
                    
                    doc = self.documents[idx].copy()
                    doc['similarity_score'] = float(score)
                    results.append(doc)
                
                return results
            except Exception as e:
                logger.error(f"Error in vector search: {e}")
        
        # Fallback to simple text matching
        return self._simple_search(query, top_k)
    
    def _simple_search(self, query: str, top_k: int) -> List[Dict[str, Any]]:
        """Simple keyword-based search fallback"""
        if not self.documents:
            return []
        
        query_lower = query.lower()
        query_words = set(query_lower.split())
        
        scored_docs = []
        for doc in self.documents:
            text_lower = doc.get('text', '').lower()
            text_words = set(text_lower.split())
            
            # Simple word overlap score
            overlap = len(query_words & text_words)
            if overlap > 0:
                score = overlap / max(len(query_words), 1)
                doc_copy = doc.copy()
                doc_copy['similarity_score'] = min(score, 1.0)
                scored_docs.append(doc_copy)
        
        # Sort by score and return top_k
        scored_docs.sort(key=lambda x: x['similarity_score'], reverse=True)
        return scored_docs[:top_k]
    
    def find_duplicates(self, query: str, threshold: float = 0.9) -> List[Dict[str, Any]]:
        """Find potential duplicate requests"""
        return self.search(query, top_k=5, threshold=threshold)


# Global instance
embeddings_manager = EmbeddingsManager()


def initialize_embeddings(documents: Optional[List[Dict[str, Any]]] = None) -> bool:
    """Initialize the global embeddings manager"""
    return embeddings_manager.initialize(documents)


def search_similar(query: str, top_k: int = 5) -> List[Dict[str, Any]]:
    """Search for similar documents"""
    return embeddings_manager.search(query, top_k)


def find_duplicates(query: str) -> List[Dict[str, Any]]:
    """Find potential duplicate requests"""
    return embeddings_manager.find_duplicates(query)
