utils/vector_storage_manager.py

"""
Vector Storage Manager - Traditional vector storage backend for dual storage comparison.
Provides vector embeddings storage with local fallback and future Modal integration.
"""

import os
import json
import time
import logging
from typing import Dict, List, Any, Optional
from pathlib import Path
import numpy as np

try:
    from sentence_transformers import SentenceTransformer
    import faiss

    VECTOR_DEPS_AVAILABLE = True
except ImportError:
    logging.warning(
        "Vector storage dependencies not available (sentence-transformers, faiss)"
    )
    SentenceTransformer = None
    faiss = None
    VECTOR_DEPS_AVAILABLE = False


class VectorStorageManager:
    """
    Vector storage backend for dual storage comparison.
    Provides traditional embedding-based storage with local FAISS index.
    Future: Modal integration for production scaling.
    """

    def __init__(
        self,
        data_dir: str = "data",
        model_name: str = "all-MiniLM-L6-v2",
        storage_handler=None,
    ):
        """
        Initialize vector storage manager.

        Args:
            data_dir (str): Base directory for storage
            model_name (str): Sentence transformer model name
            storage_handler: HF Dataset storage handler for persistence
        """
        self.logger = logging.getLogger(__name__)
        self.data_dir = Path(data_dir)
        self.model_name = model_name
        self.storage_handler = storage_handler  # For HF Dataset persistence

        # Initialize embedding model
        self.encoder = None
        if VECTOR_DEPS_AVAILABLE:
            try:
                self.encoder = SentenceTransformer(model_name)
                self.logger.info(f"Vector storage initialized with model: {model_name}")
            except Exception as e:
                self.logger.error(f"Failed to load embedding model: {e}")
        else:
            self.logger.warning("Vector storage not available - missing dependencies")

        # Client indices
        self.client_indices = {}  # client_id -> faiss index
        self.client_texts = {}  # client_id -> list of texts
        self.client_metadata = {}  # client_id -> list of metadata

    def store_embedding(
        self, text: str, client_id: str, metadata: Dict[str, Any] = None
    ) -> str:
        """
        Store text as vector embedding.

        Args:
            text (str): Text to store
            client_id (str): Client identifier
            metadata (dict): Additional metadata

        Returns:
            str: Storage result message
        """
        try:
            if not VECTOR_DEPS_AVAILABLE:
                return "Error: Vector storage dependencies not available (sentence-transformers, faiss)"

            if not self.encoder:
                return "Error: Embedding model not loaded"

            # Generate embedding
            start_time = time.time()
            embedding = self.encoder.encode([text])
            embedding_time = time.time() - start_time

            # Initialize client storage if needed
            if client_id not in self.client_indices:
                self._init_client_storage(client_id, embedding.shape[1])

            # Add to client index
            self.client_indices[client_id].add(embedding)
            self.client_texts[client_id].append(text)
            self.client_metadata[client_id].append(metadata or {})

            # Save to disk
            self._save_client_index(client_id)

            # Auto-backup to HF Dataset for persistence on HF Spaces
            self.auto_backup_after_store(client_id, self.storage_handler)

            total_embeddings = len(self.client_texts[client_id])

            return f"Vector embedding stored for client {client_id}. Embedding time: {embedding_time:.3f}s. Total embeddings: {total_embeddings}"

        except Exception as e:
            error_msg = f"Error storing vector embedding: {str(e)}"
            self.logger.error(error_msg)
            return error_msg

    def search_embeddings(self, query: str, client_id: str, top_k: int = 5) -> str:
        """
        Search embeddings using vector similarity.

        Args:
            query (str): Search query
            client_id (str): Client identifier
            top_k (int): Number of results

        Returns:
            str: JSON string with search results
        """
        try:
            if not VECTOR_DEPS_AVAILABLE:
                return json.dumps(
                    {"error": "Vector storage dependencies not available"}
                )

            if not self.encoder:
                return json.dumps({"error": "Embedding model not loaded"})

            if client_id not in self.client_indices:
                return json.dumps(
                    {"error": f"No embeddings found for client {client_id}"}
                )

            # Generate query embedding
            query_embedding = self.encoder.encode([query])

            # Search index
            scores, indices = self.client_indices[client_id].search(
                query_embedding, top_k
            )

            # Prepare results
            results = []
            for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
                if idx < len(self.client_texts[client_id]):
                    result = {
                        "text": self.client_texts[client_id][idx],
                        "score": float(score),
                        "rank": i + 1,
                        "metadata": self.client_metadata[client_id][idx],
                    }
                    results.append(result)

            return json.dumps(
                {
                    "query": query,
                    "client_id": client_id,
                    "total_results": len(results),
                    "results": results,
                    "backend": "vector_storage",
                },
                indent=2,
            )

        except Exception as e:
            error_msg = f"Error searching vector embeddings: {str(e)}"
            self.logger.error(error_msg)
            return json.dumps({"error": error_msg})

    def delete_memory(self, client_id: str, memory_name: str = "") -> str:
        """
        Delete embeddings for a client.

        Args:
            client_id (str): Client identifier
            memory_name (str): Memory name (not used in vector storage)

        Returns:
            str: Deletion result
        """
        try:
            if client_id in self.client_indices:
                # Clear client data
                del self.client_indices[client_id]
                del self.client_texts[client_id]
                del self.client_metadata[client_id]

                # Remove saved files
                client_dir = self._get_client_dir(client_id)
                if client_dir.exists():
                    import shutil

                    shutil.rmtree(client_dir)

                return f"Vector embeddings deleted for client {client_id}"
            else:
                return f"No vector embeddings found for client {client_id}"

        except Exception as e:
            error_msg = f"Error deleting vector embeddings: {str(e)}"
            self.logger.error(error_msg)
            return error_msg

    def get_stats(self, client_id: str) -> str:
        """
        Get vector storage statistics.

        Args:
            client_id (str): Client identifier

        Returns:
            str: JSON string with statistics
        """
        try:
            if client_id not in self.client_indices:
                return json.dumps(
                    {
                        "client_id": client_id,
                        "total_embeddings": 0,
                        "storage_backend": "vector_storage",
                        "status": "no_data",
                    }
                )

            total_embeddings = len(self.client_texts[client_id])
            total_text_size = sum(len(text) for text in self.client_texts[client_id])

            # Calculate storage size
            client_dir = self._get_client_dir(client_id)
            storage_size = 0
            if client_dir.exists():
                storage_size = sum(
                    f.stat().st_size for f in client_dir.rglob("*") if f.is_file()
                )

            return json.dumps(
                {
                    "client_id": client_id,
                    "total_embeddings": total_embeddings,
                    "total_text_size_bytes": total_text_size,
                    "storage_size_bytes": storage_size,
                    "storage_backend": "vector_storage",
                    "embedding_model": self.model_name,
                    "status": "active",
                },
                indent=2,
            )

        except Exception as e:
            error_msg = f"Error getting vector storage stats: {str(e)}"
            self.logger.error(error_msg)
            return json.dumps({"error": error_msg})

    def _init_client_storage(self, client_id: str, embedding_dim: int) -> None:
        """Initialize storage for a new client."""
        # Create FAISS index
        self.client_indices[client_id] = faiss.IndexFlatIP(
            embedding_dim
        )  # Inner product similarity
        self.client_texts[client_id] = []
        self.client_metadata[client_id] = []

        # Create client directory
        client_dir = self._get_client_dir(client_id)
        client_dir.mkdir(parents=True, exist_ok=True)

    def _get_client_dir(self, client_id: str) -> Path:
        """Get client-specific directory for vector storage."""
        return self.data_dir / f"{client_id}_vector"

    def _save_client_index(self, client_id: str) -> None:
        """Save client index and data to disk."""
        try:
            client_dir = self._get_client_dir(client_id)

            # Save FAISS index
            faiss.write_index(
                self.client_indices[client_id], str(client_dir / "vector_index.faiss")
            )

            # Save texts and metadata
            with open(client_dir / "texts.json", "w", encoding="utf-8") as f:
                json.dump(self.client_texts[client_id], f, indent=2)

            with open(client_dir / "metadata.json", "w", encoding="utf-8") as f:
                json.dump(self.client_metadata[client_id], f, indent=2)

        except Exception as e:
            self.logger.error(f"Error saving client index for {client_id}: {e}")

    def _load_client_index(self, client_id: str) -> bool:
        """Load client index and data from disk."""
        try:
            client_dir = self._get_client_dir(client_id)

            if not (client_dir / "vector_index.faiss").exists():
                return False

            # Load FAISS index
            self.client_indices[client_id] = faiss.read_index(
                str(client_dir / "vector_index.faiss")
            )

            # Load texts and metadata
            with open(client_dir / "texts.json", "r", encoding="utf-8") as f:
                self.client_texts[client_id] = json.load(f)

            with open(client_dir / "metadata.json", "r", encoding="utf-8") as f:
                self.client_metadata[client_id] = json.load(f)

            return True

        except Exception as e:
            self.logger.error(f"Error loading client index for {client_id}: {e}")
            return False

    def load_client_data(self, client_id: str) -> str:
        """
        Load client data from disk.

        Args:
            client_id (str): Client identifier

        Returns:
            str: Load result message
        """
        try:
            if self._load_client_index(client_id):
                total_embeddings = len(self.client_texts[client_id])
                return f"Vector storage loaded for client {client_id}: {total_embeddings} embeddings"
            else:
                return f"No vector storage data found for client {client_id}"

        except Exception as e:
            error_msg = f"Error loading client data: {str(e)}"
            self.logger.error(error_msg)
            return error_msg

    # Future Modal integration methods (placeholders)

    def enable_modal_backend(self, modal_token: str) -> str:
        """
        Enable Modal backend for production scaling.

        Args:
            modal_token (str): Modal API token

        Returns:
            str: Activation result
        """
        # TODO: Implement Modal integration
        return (
            "Modal backend integration not yet implemented. Using local FAISS storage."
        )

    def migrate_to_modal(self, client_id: str) -> str:
        """
        Migrate client data to Modal backend.

        Args:
            client_id (str): Client identifier

        Returns:
            str: Migration result
        """
        # TODO: Implement Modal migration
        return "Modal migration not yet implemented. Data remains in local storage."

    # HF Dataset Integration for Persistence on HF Spaces

    def backup_to_hf_dataset(self, client_id: str, storage_handler) -> str:
        """
        Backup vector storage to HuggingFace Dataset for persistence.

        Args:
            client_id (str): Client identifier
            storage_handler: HF Dataset storage handler

        Returns:
            str: Backup result
        """
        try:
            if not storage_handler or not storage_handler.hf_enabled:
                return "HF Dataset backup not available - no storage handler or HF not enabled"

            client_dir = self._get_client_dir(client_id)
            if not client_dir.exists():
                return f"No vector data found for client {client_id}"

            # Use storage handler to backup vector files
            success = storage_handler.backup_client_data(client_id, client_dir)

            if success:
                return f"Successfully backed up vector storage for client {client_id} to HF Dataset"
            else:
                return f"Failed to backup vector storage for client {client_id}"

        except Exception as e:
            error_msg = f"Error backing up vector storage: {str(e)}"
            self.logger.error(error_msg)
            return error_msg

    def restore_from_hf_dataset(self, client_id: str, storage_handler) -> str:
        """
        Restore vector storage from HuggingFace Dataset.

        Args:
            client_id (str): Client identifier
            storage_handler: HF Dataset storage handler

        Returns:
            str: Restore result
        """
        try:
            if not storage_handler or not storage_handler.hf_enabled:
                return "HF Dataset restore not available - no storage handler or HF not enabled"

            client_dir = self._get_client_dir(client_id)

            # Use storage handler to restore vector files
            success = storage_handler.restore_client_data(client_id, client_dir)

            if success:
                # Load the restored data into memory
                if self._load_client_index(client_id):
                    total_embeddings = len(self.client_texts[client_id])
                    return f"Successfully restored vector storage for client {client_id}: {total_embeddings} embeddings"
                else:
                    return f"Vector files restored but failed to load into memory for client {client_id}"
            else:
                return f"Failed to restore vector storage for client {client_id}"

        except Exception as e:
            error_msg = f"Error restoring vector storage: {str(e)}"
            self.logger.error(error_msg)
            return error_msg

    def auto_backup_after_store(self, client_id: str, storage_handler) -> None:
        """
        Automatically backup after storing embeddings (for HF Spaces persistence).

        Args:
            client_id (str): Client identifier
            storage_handler: HF Dataset storage handler
        """
        try:
            if storage_handler and storage_handler.hf_enabled:
                # Auto-backup in background (non-blocking)
                self.backup_to_hf_dataset(client_id, storage_handler)
        except Exception as e:
            self.logger.warning(f"Auto-backup failed for client {client_id}: {e}")