"""
Vector Store Abstraction Layer
Supports multiple vector databases (Chroma, FAISS, etc.)
"""

from abc import ABC, abstractmethod
from typing import List, Dict, Optional, Any
import logging

logger = logging.getLogger(__name__)


class VectorStore(ABC):
    """Abstract interface for vector stores"""
    
    @abstractmethod
    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
        """
        Add documents to vector store
        
        Args:
            documents: List of dicts with 'id', 'text', 'metadata'
        """
        pass
    
    @abstractmethod
    def search(self, query: str, limit: int = 5, filter_metadata: Optional[Dict] = None) -> List[Dict[str, Any]]:
        """
        Search similar documents
        
        Args:
            query: Search query text
            limit: Maximum number of results
            filter_metadata: Optional metadata filters
            
        Returns:
            List of similar documents with scores
        """
        pass
    
    @abstractmethod
    def delete_collection(self) -> None:
        """Delete all vectors in collection"""
        pass
    
    @abstractmethod
    def get_stats(self) -> Dict[str, Any]:
        """Get statistics about the vector store"""
        pass


try:
    import chromadb
    from chromadb.config import Settings
    CHROMA_AVAILABLE = True
except ImportError:
    CHROMA_AVAILABLE = False
    logger.warning("ChromaDB not available. Install with: pip install chromadb")


class ChromaVectorStore(VectorStore):
    """ChromaDB implementation of VectorStore"""
    
    def __init__(self, collection_name: str = "knowledge_base", 
                 persist_directory: Optional[str] = None,
                 embedding_model: str = "all-MiniLM-L6-v2"):
        """
        Initialize ChromaDB vector store
        
        Args:
            collection_name: Name of the collection
            persist_directory: Directory to persist data (None = in-memory)
            embedding_model: Embedding model name (sentence-transformers compatible)
        """
        if not CHROMA_AVAILABLE:
            raise ImportError(
                "ChromaDB is not installed. Install with: pip install chromadb"
            )
        
        self.collection_name = collection_name
        self.persist_directory = persist_directory
        self.embedding_model = embedding_model
        
        # Initialize Chroma client
        if persist_directory:
            self.client = chromadb.PersistentClient(path=persist_directory)
        else:
            self.client = chromadb.Client()
        
        # Lazy load embedding model
        self._embedding_fn = None
        
        # Get or create collection with embedding function
        try:
            # Create embedding function
            embedding_fn = self._get_embedding_function()
            
            self.collection = self.client.get_or_create_collection(
                name=collection_name,
                embedding_function=embedding_fn,
                metadata={"hnsw:space": "cosine"}
            )
        except Exception as e:
            logger.error(f"Failed to create Chroma collection: {e}")
            raise
    
    def _get_embedding_function(self):
        """Lazy load embedding function"""
        if self._embedding_fn is None:
            try:
                # Try to use ChromaDB's native SentenceTransformerEmbeddingFunction
                try:
                    # Try different import paths for ChromaDB embedding functions
                    try:
                        from chromadb.utils import embedding_functions
                        embedding_fn_class = embedding_functions.SentenceTransformerEmbeddingFunction
                    except (ImportError, AttributeError):
                        try:
                            from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction as embedding_fn_class
                        except ImportError:
                            embedding_fn_class = None
                    
                    if embedding_fn_class:
                        self._embedding_fn = embedding_fn_class(model_name=self.embedding_model)
                        logger.info(f"Loaded embedding model using ChromaDB native function: {self.embedding_model}")
                    else:
                        raise AttributeError("SentenceTransformerEmbeddingFunction not found")
                        
                except (ImportError, AttributeError, Exception) as e:
                    # Fallback: Custom embedding function wrapper compatible with ChromaDB
                    from sentence_transformers import SentenceTransformer
                    model = SentenceTransformer(self.embedding_model)
                    
                    class CustomEmbeddingFunction:
                        def __init__(self, model, model_name):
                            self.model = model
                            self.model_name = model_name
                            self.name = model_name  # ChromaDB may check for 'name' attribute
                        
                        def __call__(self, texts: List[str]) -> List[List[float]]:
                            embeddings = self.model.encode(texts, show_progress_bar=False)
                            return embeddings.tolist()
                        
                        def encode_queries(self, queries: List[str]) -> List[List[float]]:
                            return self.__call__(queries)
                    
                    self._embedding_fn = CustomEmbeddingFunction(model, self.embedding_model)
                    logger.info(f"Loaded embedding model using custom wrapper: {self.embedding_model} (fallback: {e})")
            except ImportError:
                raise ImportError(
                    "sentence-transformers not installed. "
                    "Install with: pip install sentence-transformers"
                )
        
        return self._embedding_fn
    
    def add_documents(self, documents: List[Dict[str, Any]]) -> None:
        """Add documents to ChromaDB"""
        if not documents:
            return
        
        # Prepare data
        ids = []
        texts = []
        metadatas = []
        
        for doc in documents:
            doc_id = str(doc.get('id', doc.get('text', ''))[:100])
            # Ensure unique IDs
            if doc_id in ids:
                doc_id = f"{doc_id}_{len(ids)}"
            ids.append(doc_id)
            texts.append(doc['text'])
            # Ensure metadata values are JSON-serializable
            metadata = doc.get('metadata', {})
            clean_metadata = {}
            for k, v in metadata.items():
                if isinstance(v, (str, int, float, bool)) or v is None:
                    clean_metadata[k] = v
                else:
                    clean_metadata[k] = str(v)
            metadatas.append(clean_metadata)
        
        # Add to collection (Chroma will use embedding function automatically)
        try:
            self.collection.add(
                ids=ids,
                documents=texts,
                metadatas=metadatas
            )
            
            logger.debug(f"Added {len(documents)} documents to Chroma")
        except Exception as e:
            logger.error(f"Error adding documents to Chroma: {e}")
            raise
    
    def search(self, query: str, limit: int = 5, 
               filter_metadata: Optional[Dict] = None) -> List[Dict[str, Any]]:
        """Search in ChromaDB"""
        try:
            # Build where clause for metadata filtering
            where = None
            if filter_metadata:
                where = filter_metadata
            
            # Search (Chroma will use embedding function automatically)
            results = self.collection.query(
                query_texts=[query],
                n_results=limit,
                where=where
            )
            
            # Format results
            formatted_results = []
            if results.get('documents') and len(results['documents']) > 0 and len(results['documents'][0]) > 0:
                num_results = len(results['documents'][0])
                distances = results.get('distances', [[0.0] * num_results])
                
                for i in range(num_results):
                    # ChromaDB uses cosine distance (0 = identical, 1 = opposite)
                    # Convert to similarity score (1 = identical, 0 = opposite)
                    distance = distances[0][i] if distances and len(distances[0]) > i else 0.0
                    similarity = 1.0 - distance if distance <= 1.0 else max(0.0, 1.0 / (1.0 + distance))
                    
                    formatted_results.append({
                        'id': results['ids'][0][i] if results.get('ids') and len(results['ids'][0]) > i else f"doc_{i}",
                        'text': results['documents'][0][i],
                        'metadata': results['metadatas'][0][i] if results.get('metadatas') and len(results['metadatas'][0]) > i else {},
                        'score': similarity
                    })
            
            return formatted_results
        except Exception as e:
            logger.error(f"Error searching Chroma: {e}")
            return []
    
    def delete_collection(self) -> None:
        """Delete collection"""
        try:
            self.client.delete_collection(self.collection_name)
            logger.info(f"Deleted Chroma collection: {self.collection_name}")
        except Exception as e:
            logger.error(f"Error deleting collection: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get collection statistics"""
        try:
            count = self.collection.count()
            return {
                'total_documents': count,
                'collection_name': self.collection_name,
                'embedding_model': self.embedding_model
            }
        except Exception as e:
            logger.error(f"Error getting stats: {e}")
            return {'total_documents': 0}


def create_vector_store(store_type: str = "chroma", **kwargs) -> Optional[VectorStore]:
    """
    Factory function to create vector store
    
    Args:
        store_type: Type of vector store ('chroma', 'faiss', etc.)
        **kwargs: Store-specific parameters
        
    Returns:
        VectorStore instance or None if not available
    """
    if store_type == "chroma":
        if not CHROMA_AVAILABLE:
            logger.warning("ChromaDB not available. Install with: pip install chromadb")
            return None
        return ChromaVectorStore(**kwargs)
    else:
        logger.warning(f"Unknown vector store type: {store_type}")
        return None

