#!/usr/bin/env python3
"""
Pre-load sentence-transformers models to avoid startup delays.

This script downloads and caches the default embedding models used by
MCP Memory Service so they don't need to be downloaded during server startup,
which can cause timeout errors in Claude Desktop.
"""

import sys
import os

def preload_sentence_transformers():
    """Pre-load the default sentence-transformers model."""
    try:
        print("[INFO] Pre-loading sentence-transformers models...")
        from sentence_transformers import SentenceTransformer
        
        # Default model used by the memory service
        model_name = "all-MiniLM-L6-v2"
        print(f"[INFO] Downloading and caching model: {model_name}")
        
        model = SentenceTransformer(model_name)
        print(f"[OK] Model loaded successfully on device: {model.device}")
        
        # Test the model to ensure it works
        print("[INFO] Testing model functionality...")
        test_text = "This is a test sentence for embedding."
        embedding = model.encode(test_text)
        print(f"[OK] Model test successful - embedding shape: {embedding.shape}")
        
        return True
        
    except ImportError:
        print("[WARNING] sentence-transformers not available - skipping model preload")
        return False
    except Exception as e:
        print(f"[ERROR] Error preloading model: {e}")
        return False

def check_cache_status():
    """Check if models are already cached."""
    cache_locations = [
        os.path.expanduser("~/.cache/huggingface/hub"),
        os.path.expanduser("~/.cache/torch/sentence_transformers"),
    ]
    
    for cache_dir in cache_locations:
        if os.path.exists(cache_dir):
            try:
                contents = os.listdir(cache_dir)
                for item in contents:
                    if 'sentence-transformers' in item.lower() or 'minilm' in item.lower():
                        print(f"[OK] Found cached model: {item}")
                        return True
            except (OSError, PermissionError):
                continue
    
    print("[INFO] No cached models found")
    return False

def main():
    print("MCP Memory Service - Model Preloader")
    print("=" * 50)
    
    # Check current cache status
    print("\n[1] Checking cache status...")
    models_cached = check_cache_status()
    
    if models_cached:
        print("[OK] Models are already cached - no download needed")
        return True
    
    # Preload models
    print("\n[2] Preloading models...")
    success = preload_sentence_transformers()
    
    if success:
        print("\n[OK] Model preloading complete!")
        print("[INFO] MCP Memory Service should now start without downloading models")
    else:
        print("\n[WARNING] Model preloading failed - server may need to download during startup")
        
    return success

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)