MCP-MCP

"""Semantic search functionality for MCP server discovery using sentence transformers."""

import hashlib
import json
import os
from pathlib import Path
from typing import Optional

import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity

from .database import MCPServerEntry
from settings import app_logger

logger = app_logger.getChild(__name__)

# Model configuration
DEFAULT_MODEL = "all-MiniLM-L6-v2"  # Fast, lightweight, good quality
EMBEDDINGS_VERSION = "v1"  # Increment when changing embedding logic


def get_cache_dir() -> Path:
    """Get cache directory following XDG Base Directory Specification."""
    # Try XDG_CACHE_HOME first, fallback to ~/.cache
    xdg_cache_home = os.environ.get("XDG_CACHE_HOME")
    if xdg_cache_home:
        cache_base = Path(xdg_cache_home)
    else:
        cache_base = Path.home() / ".cache"
    
    cache_dir = cache_base / "mcp-mcp" / "embeddings"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


class SemanticSearchEngine:
    """Semantic search engine for MCP servers using sentence transformers."""
    
    def __init__(self, model_name: str = DEFAULT_MODEL):
        self.model_name = model_name
        self.model: Optional[SentenceTransformer] = None
        self.server_embeddings: Optional[np.ndarray] = None
        self.servers: list[MCPServerEntry] = []
        self.embeddings_hash: Optional[str] = None
        
        # Cache directory will be created by get_cache_dir()
    
    def _get_server_texts(self, servers: list[MCPServerEntry]) -> list[str]:
        """Extract text content from servers for embedding."""
        texts = []
        for server in servers:
            # Combine name, description, and category for richer context
            text = f"{server.name}. {server.description}. Category: {server.category}"
            texts.append(text)
        return texts
    
    def _compute_content_hash(self, servers: list[MCPServerEntry]) -> str:
        """Compute hash of server content for cache invalidation."""
        content = []
        for server in servers:
            content.append({
                "name": server.name,
                "description": server.description,
                "category": server.category,
                "url": server.url
            })
        
        # Sort for consistent hashing
        content.sort(key=lambda x: x["name"])
        content_str = json.dumps(content, sort_keys=True)
        
        # Include model name and version in hash
        hash_input = f"{EMBEDDINGS_VERSION}:{self.model_name}:{content_str}"
        return hashlib.sha256(hash_input.encode()).hexdigest()[:16]
    
    def _get_cache_path(self, content_hash: str) -> Path:
        """Get path for cached embeddings file."""
        return get_cache_dir() / f"embeddings_{content_hash}.npy"
    
    def _load_cached_embeddings(self, content_hash: str) -> Optional[np.ndarray]:
        """Load embeddings from cache if available."""
        cache_path = self._get_cache_path(content_hash)
        if cache_path.exists():
            try:
                embeddings = np.load(cache_path)
                logger.debug(f"Loaded cached embeddings from {cache_path}")
                return embeddings
            except Exception as e:
                logger.warning(f"Failed to load cached embeddings: {e}")
                # Delete corrupted cache file
                cache_path.unlink(missing_ok=True)
        return None
    
    def _save_embeddings_to_cache(self, embeddings: np.ndarray, content_hash: str):
        """Save embeddings to cache."""
        cache_path = self._get_cache_path(content_hash)
        try:
            np.save(cache_path, embeddings)
            logger.debug(f"Saved embeddings to cache: {cache_path}")
        except Exception as e:
            logger.warning(f"Failed to save embeddings to cache: {e}")
    
    def _cleanup_old_cache_files(self):
        """Remove old cache files to prevent disk bloat."""
        try:
            cache_dir = get_cache_dir()
            cache_files = list(cache_dir.glob("embeddings_*.npy"))
            # Keep only the most recent 5 cache files
            if len(cache_files) > 5:
                cache_files.sort(key=lambda p: p.stat().st_mtime)
                for old_file in cache_files[:-5]:
                    old_file.unlink()
                    logger.debug(f"Removed old cache file: {old_file}")
        except Exception as e:
            logger.warning(f"Failed to cleanup cache files: {e}")
    
    async def initialize(self, servers: list[MCPServerEntry]):
        """Initialize the semantic search engine with server data."""
        self.servers = servers
        content_hash = self._compute_content_hash(servers)
        self.embeddings_hash = content_hash
        
        # Try to load from cache first
        cached_embeddings = self._load_cached_embeddings(content_hash)
        if cached_embeddings is not None:
            self.server_embeddings = cached_embeddings
            logger.info(f"Using cached embeddings for {len(servers)} servers")
            return
        
        # Load model if not already loaded
        if self.model is None:
            logger.info(f"Loading sentence transformer model: {self.model_name}")
            try:
                self.model = SentenceTransformer(self.model_name)
                logger.info("Model loaded successfully")
            except Exception as e:
                logger.error(f"Failed to load model {self.model_name}: {e}")
                raise
        
        # Generate embeddings
        logger.info(f"Generating embeddings for {len(servers)} servers...")
        texts = self._get_server_texts(servers)
        
        try:
            # Generate embeddings in batch for efficiency
            self.server_embeddings = self.model.encode(
                texts,
                convert_to_numpy=True,
                show_progress_bar=True if len(texts) > 10 else False
            )
            logger.info("Embeddings generated successfully")
            
            # Save to cache
            self._save_embeddings_to_cache(self.server_embeddings, content_hash)
            
            # Cleanup old cache files
            self._cleanup_old_cache_files()
            
        except Exception as e:
            logger.error(f"Failed to generate embeddings: {e}")
            raise
    
    async def initialize_with_precomputed_embeddings(
        self, 
        servers: list[MCPServerEntry], 
        precomputed_embeddings: np.ndarray
    ):
        """Initialize the semantic search engine with precomputed embeddings."""
        self.servers = servers
        self.server_embeddings = precomputed_embeddings
        
        # Validate that embeddings match server count
        if len(servers) != precomputed_embeddings.shape[0]:
            raise ValueError(
                f"Server count ({len(servers)}) does not match embeddings count "
                f"({precomputed_embeddings.shape[0]})"
            )
        
        # Load model for query encoding (still needed for search)
        if self.model is None:
            logger.info(f"Loading sentence transformer model: {self.model_name}")
            try:
                self.model = SentenceTransformer(self.model_name)
                logger.info("Model loaded successfully")
            except Exception as e:
                logger.error(f"Failed to load model {self.model_name}: {e}")
                raise
        
        # Generate a simple hash for the precomputed data
        content_hash = f"precomputed_{len(servers)}_{precomputed_embeddings.shape[1]}"
        self.embeddings_hash = content_hash
        
        logger.info(f"Initialized with precomputed embeddings for {len(servers)} servers")
    
    def semantic_search(
        self, 
        query: str, 
        top_k: int = 10,
        similarity_threshold: float = 0.1
    ) -> list[tuple[MCPServerEntry, float]]:
        """
        Perform semantic search for MCP servers.
        
        Args:
            query: Search query
            top_k: Maximum number of results to return
            similarity_threshold: Minimum similarity score (0-1)
            
        Returns:
            List of (server, similarity_score) tuples, sorted by similarity
        """
        if self.model is None:
            raise RuntimeError("Model not loaded. Call initialize() first.")
        
        if self.server_embeddings is None:
            raise RuntimeError("Server embeddings not available. Call initialize() first.")
        
        # Generate query embedding
        try:
            query_embedding = self.model.encode([query], convert_to_numpy=True)
        except Exception as e:
            logger.error(f"Failed to encode query: {e}")
            return []
        
        # Compute similarities
        try:
            similarities = cosine_similarity(query_embedding, self.server_embeddings)[0]
        except Exception as e:
            logger.error(f"Failed to compute similarities: {e}")
            return []
        
        # Create results with similarity scores
        results = []
        for i, similarity in enumerate(similarities):
            if similarity >= similarity_threshold:
                results.append((self.servers[i], float(similarity)))
        
        # Sort by similarity (highest first) and limit results
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]
    
    def is_available(self) -> bool:
        """Check if semantic search is available (model loaded and embeddings ready)."""
        return (
            self.model is not None and 
            self.server_embeddings is not None and 
            len(self.servers) > 0
        )
    
    def get_cache_info(self) -> dict:
        """Get information about the current cache state."""
        cache_dir = get_cache_dir()
        cache_files = list(cache_dir.glob("embeddings_*.npy"))
        current_cache = None
        if self.embeddings_hash:
            current_cache = self._get_cache_path(self.embeddings_hash)
        
        return {
            "cache_dir": str(cache_dir),
            "total_cache_files": len(cache_files),
            "current_cache_file": str(current_cache) if current_cache else None,
            "current_cache_exists": current_cache.exists() if current_cache else False,
            "model_loaded": self.model is not None,
            "embeddings_ready": self.server_embeddings is not None,
            "num_servers": len(self.servers)
        }