embeddings.py

"""
Embedding service for semantic search using ChromaDB and sentence-transformers.

Manages persistent vector store and provides search operations.
CSV remains source of truth; ChromaDB is a derived index.
"""
from pathlib import Path
from typing import Optional, List, Dict, Any, Callable
import logging

# Data paths (same pattern as operations.py)
DATA_DIR = Path(__file__).parent / "data"
CHROMA_DIR = DATA_DIR / "chroma_db"
MEMORY_MATTER_CSV = DATA_DIR / "memory_matter.csv"


class EmbeddingService:
    """
    Singleton embedding service for semantic memory search.

    Features:
    - Persistent ChromaDB storage in ./data/chroma_db/
    - Lazy model loading (first use)
    - Metadata filtering on search
    - Sync status tracking
    """

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        collection_name: str = "memory_matter"
    ):
        self.model_name = model_name
        self.collection_name = collection_name
        self._model = None
        self._client = None
        self._collection = None
        self._initialized = False
        self._init_error = None

    def _lazy_init(self) -> bool:
        """Initialize ChromaDB and model on first use. Returns success status."""
        if self._initialized:
            return True

        if self._init_error:
            return False

        try:
            import chromadb
            from sentence_transformers import SentenceTransformer

            # Initialize persistent client
            CHROMA_DIR.mkdir(parents=True, exist_ok=True)
            self._client = chromadb.PersistentClient(path=str(CHROMA_DIR))

            # Get or create collection with cosine similarity
            self._collection = self._client.get_or_create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"}
            )

            # Load embedding model
            self._model = SentenceTransformer(self.model_name)

            self._initialized = True
            return True

        except ImportError as e:
            self._init_error = f"Missing dependencies: {e}"
            logging.warning(f"[embeddings] {self._init_error}")
            return False
        except Exception as e:
            self._init_error = str(e)
            logging.warning(f"[embeddings] Init failed: {self._init_error}")
            return False

    def is_available(self) -> bool:
        """Check if embeddings service is available and initialized."""
        if self._initialized:
            return True
        return self._lazy_init()

    def embed_text(self, text: str) -> List[float]:
        """Generate embedding vector for text."""
        if not self._lazy_init():
            raise RuntimeError(f"Embedding service unavailable: {self._init_error}")
        return self._model.encode(text).tolist()

    def add_memory(
        self,
        memory_id: str,
        raw: str,
        metadata: Dict[str, Any]
    ) -> bool:
        """
        Add or update a memory embedding.

        Args:
            memory_id: Memory ID (e.g., mm_001)
            raw: The raw memory text to embed
            metadata: Additional metadata (timestamp, layers, tags, etc.)
        """
        if not self._lazy_init():
            return False

        try:
            # Generate embedding
            embedding = self.embed_text(raw)

            # ChromaDB metadata must be string/int/float/bool
            clean_metadata = {}
            for k, v in metadata.items():
                if v is None:
                    continue
                if isinstance(v, (str, int, float, bool)):
                    clean_metadata[k] = v
                else:
                    clean_metadata[k] = str(v)

            # Upsert (add or update)
            self._collection.upsert(
                ids=[memory_id],
                embeddings=[embedding],
                documents=[raw],
                metadatas=[clean_metadata]
            )
            return True

        except Exception as e:
            logging.warning(f"[embeddings] Failed to add {memory_id}: {e}")
            return False

    def remove_memory(self, memory_id: str) -> bool:
        """Remove a memory from the index."""
        if not self._lazy_init():
            return False

        try:
            self._collection.delete(ids=[memory_id])
            return True
        except Exception as e:
            logging.warning(f"[embeddings] Failed to remove {memory_id}: {e}")
            return False

    def semantic_search(
        self,
        query: str,
        limit: int = 10,
        layer: Optional[str] = None,
        min_toxicity: Optional[int] = None,
        max_toxicity: Optional[int] = None,
        min_nutrient: Optional[int] = None,
        max_nutrient: Optional[int] = None,
        date_start: Optional[str] = None,
        date_end: Optional[str] = None
    ) -> List[Dict[str, Any]]:
        """
        Semantic search with optional metadata filters.

        Args:
            query: Natural language search query
            limit: Max results to return
            layer: Filter by layer (substring match in layers field)
            min_toxicity, max_toxicity: Toxicity range filter
            min_nutrient, max_nutrient: Nutrient density range filter
            date_start, date_end: Date range filter (YYYY-MM-DD strings)

        Returns:
            List of {id, score, raw, metadata} dicts sorted by relevance
        """
        if not self._lazy_init():
            return []

        try:
            # Build where clause for filters
            where_conditions = []

            # Note: ChromaDB uses exact match for strings, not substring
            # Layer filtering is limited - will do post-filter if needed
            if layer:
                where_conditions.append({"layers": {"$contains": layer}})

            if min_toxicity is not None:
                where_conditions.append({"toxicity": {"$gte": min_toxicity}})
            if max_toxicity is not None:
                where_conditions.append({"toxicity": {"$lte": max_toxicity}})
            if min_nutrient is not None:
                where_conditions.append({"nutrient_density": {"$gte": min_nutrient}})
            if max_nutrient is not None:
                where_conditions.append({"nutrient_density": {"$lte": max_nutrient}})
            if date_start:
                where_conditions.append({"timestamp_actual": {"$gte": date_start}})
            if date_end:
                where_conditions.append({"timestamp_actual": {"$lte": date_end}})

            # Combine conditions
            where = None
            if len(where_conditions) == 1:
                where = where_conditions[0]
            elif len(where_conditions) > 1:
                where = {"$and": where_conditions}

            # Generate query embedding
            query_embedding = self.embed_text(query)

            # Search with more results to account for filtering
            search_limit = limit * 2 if where else limit

            results = self._collection.query(
                query_embeddings=[query_embedding],
                n_results=search_limit,
                where=where,
                include=["documents", "metadatas", "distances"]
            )

            # Format results
            formatted = []
            if results and results['ids'] and results['ids'][0]:
                for i, memory_id in enumerate(results['ids'][0]):
                    # Convert distance to similarity score (cosine distance to similarity)
                    distance = results['distances'][0][i]
                    score = 1 - distance  # Cosine similarity

                    metadata = results['metadatas'][0][i] if results['metadatas'] else {}

                    # Post-filter for layer substring match if needed
                    if layer and 'layers' in metadata:
                        if layer.lower() not in metadata['layers'].lower():
                            continue

                    formatted.append({
                        'id': memory_id,
                        'score': round(score, 4),
                        'raw': results['documents'][0][i],
                        'metadata': metadata
                    })

                    if len(formatted) >= limit:
                        break

            return formatted

        except Exception as e:
            logging.warning(f"[embeddings] Search failed: {e}")
            return []

    def find_similar(
        self,
        memory_id: str,
        limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find memories similar to a given memory.

        Args:
            memory_id: Source memory ID
            limit: Max similar memories to return

        Returns:
            List of similar memories (excluding the source)
        """
        if not self._lazy_init():
            return []

        try:
            # Get the memory's embedding
            result = self._collection.get(
                ids=[memory_id],
                include=["embeddings"]
            )

            if not result or result.get('embeddings') is None or len(result['embeddings']) == 0:
                return []

            embedding = result['embeddings'][0]
            if embedding is None:
                return []

            # Search, getting extra to exclude self
            results = self._collection.query(
                query_embeddings=[embedding],
                n_results=limit + 1,
                include=["documents", "metadatas", "distances"]
            )

            # Format and exclude self
            formatted = []
            if results and results['ids'] and results['ids'][0]:
                for i, mid in enumerate(results['ids'][0]):
                    if mid == memory_id:
                        continue

                    distance = results['distances'][0][i]
                    score = 1 - distance

                    formatted.append({
                        'id': mid,
                        'score': round(score, 4),
                        'raw': results['documents'][0][i],
                        'metadata': results['metadatas'][0][i] if results['metadatas'] else {}
                    })

                    if len(formatted) >= limit:
                        break

            return formatted

        except Exception as e:
            logging.warning(f"[embeddings] Find similar failed: {e}")
            return []

    def get_sync_status(self) -> Dict[str, Any]:
        """
        Get sync status between CSV and ChromaDB.

        Returns dict with: initialized, indexed_count, csv_count, in_sync
        """
        import pandas as pd

        # Count CSV entries
        csv_count = 0
        if MEMORY_MATTER_CSV.exists():
            try:
                df = pd.read_csv(MEMORY_MATTER_CSV)
                csv_count = len(df)
            except Exception:
                pass

        if not self._initialized and not self._lazy_init():
            return {
                'initialized': False,
                'indexed_count': 0,
                'csv_count': csv_count,
                'in_sync': False,
                'error': self._init_error
            }

        try:
            indexed_count = self._collection.count()

            return {
                'initialized': True,
                'indexed_count': indexed_count,
                'csv_count': csv_count,
                'in_sync': indexed_count == csv_count
            }
        except Exception as e:
            return {
                'initialized': False,
                'indexed_count': 0,
                'csv_count': csv_count,
                'in_sync': False,
                'error': str(e)
            }

    def rebuild_index(
        self,
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> Dict[str, Any]:
        """
        Rebuild entire index from CSV.

        Args:
            progress_callback: Optional callback(current, total) for progress updates

        Returns:
            Dict with success, total, indexed, errors
        """
        import pandas as pd

        if not self._lazy_init():
            return {
                'success': False,
                'total': 0,
                'indexed': 0,
                'errors': [{'error': self._init_error}]
            }

        try:
            # Clear existing collection
            try:
                self._client.delete_collection(self.collection_name)
            except Exception:
                pass

            # Recreate collection
            self._collection = self._client.create_collection(
                name=self.collection_name,
                metadata={"hnsw:space": "cosine"}
            )

            # Read all memories from CSV
            if not MEMORY_MATTER_CSV.exists():
                return {
                    'success': True,
                    'total': 0,
                    'indexed': 0,
                    'errors': []
                }

            df = pd.read_csv(MEMORY_MATTER_CSV)
            total = len(df)
            indexed = 0
            errors = []

            for i, row in df.iterrows():
                try:
                    memory_id = row['id']
                    raw = str(row.get('raw', ''))

                    if not raw or raw == 'nan':
                        continue

                    metadata = {
                        'timestamp_actual': str(row.get('timestamp_actual', '')),
                        'layers': str(row.get('layers', '')),
                        'tags': str(row.get('tags', '')),
                        'source_type': str(row.get('source_type', ''))
                    }

                    # Handle numeric fields
                    try:
                        metadata['toxicity'] = int(row.get('toxicity', 0))
                    except (ValueError, TypeError):
                        metadata['toxicity'] = 0

                    try:
                        metadata['nutrient_density'] = int(row.get('nutrient_density', 0))
                    except (ValueError, TypeError):
                        metadata['nutrient_density'] = 0

                    self.add_memory(memory_id, raw, metadata)
                    indexed += 1

                    if progress_callback:
                        progress_callback(i + 1, total)

                except Exception as e:
                    errors.append({'id': str(row.get('id', f'row_{i}')), 'error': str(e)})

            return {
                'success': True,
                'total': total,
                'indexed': indexed,
                'errors': errors
            }

        except Exception as e:
            return {
                'success': False,
                'total': 0,
                'indexed': 0,
                'errors': [{'error': str(e)}]
            }


# Singleton instance
embeddings = EmbeddingService()