#!/usr/bin/env python3
"""
Test mem0 configuration validation with Supabase
"""

import os
import sys
from dotenv import load_dotenv
from config import load_config, get_mem0_config
import vecs

def test_supabase_vector_store_connection():
    """Test direct connection to Supabase vector store using vecs"""
    print("🔗 Testing direct Supabase vector store connection...")
    
    try:
        # Connection string for our Supabase instance
        connection_string = "postgresql://supabase_admin:CzkaYmRvc26Y@localhost:5435/postgres"
        
        # Create vecs client
        db = vecs.create_client(connection_string)
        
        # List existing collections
        collections = db.list_collections()
        print(f"✅ Connected to Supabase PostgreSQL")
        print(f"📦 Existing collections: {[c.name for c in collections]}")
        
        # Test creating a collection (this will create the table if it doesn't exist)
        collection_name = "mem0_test_vectors"
        collection = db.get_or_create_collection(
            name=collection_name,
            dimension=1536  # OpenAI text-embedding-3-small dimension
        )
        
        print(f"✅ Collection '{collection_name}' ready")
        
        # Test basic vector operations
        print("🧪 Testing basic vector operations...")
        
        # Insert a test vector
        test_id = "test_vector_1"
        test_vector = [0.1] * 1536  # Dummy vector
        test_metadata = {"content": "This is a test memory", "user_id": "test_user"}
        
        collection.upsert(
            records=[(test_id, test_vector, test_metadata)]
        )
        print("✅ Vector upserted successfully")
        
        # Search for similar vectors
        query_vector = [0.1] * 1536  # Same as test vector
        results = collection.query(
            data=query_vector,
            limit=5,
            include_metadata=True
        )
        
        print(f"✅ Search completed, found {len(results)} results")
        if results:
            print(f"   First result: {results[0]}")
        
        # Cleanup
        collection.delete(ids=[test_id])
        print("✅ Test data cleaned up")
        
        return True
        
    except Exception as e:
        print(f"❌ Supabase vector store connection failed: {e}")
        import traceback
        traceback.print_exc()
        return False

def test_configuration_validation():
    """Test mem0 configuration validation"""
    print("⚙️  Testing mem0 configuration validation...")
    
    try:
        config = load_config()
        mem0_config = get_mem0_config(config, "openai")
        
        print("✅ Configuration loaded successfully")
        print(f"📋 Vector store provider: {mem0_config['vector_store']['provider']}")
        print(f"📋 Graph store provider: {mem0_config.get('graph_store', {}).get('provider', 'Not configured')}")
        
        # Validate required fields
        vector_config = mem0_config['vector_store']['config']
        required_fields = ['connection_string', 'collection_name', 'embedding_model_dims']
        
        for field in required_fields:
            if field not in vector_config:
                raise ValueError(f"Missing required field: {field}")
        
        print("✅ All required configuration fields present")
        return True
        
    except Exception as e:
        print(f"❌ Configuration validation failed: {e}")
        return False

def main():
    """Main test function"""
    print("=" * 60)
    print("MEM0 + SUPABASE CONFIGURATION TESTS")
    print("=" * 60)
    
    # Load environment
    load_dotenv()
    
    results = []
    
    # Test 1: Configuration validation
    results.append(("Configuration Validation", test_configuration_validation()))
    
    # Test 2: Direct Supabase vector store connection
    results.append(("Supabase Vector Store", test_supabase_vector_store_connection()))
    
    # Summary
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    passed = 0
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} {test_name}")
        if result:
            passed += 1
    
    print(f"\nOverall: {passed}/{len(results)} tests passed")
    
    if passed == len(results):
        print("🎉 Supabase configuration is ready!")
        sys.exit(0)
    else:
        print("💥 Some tests failed - check configuration")
        sys.exit(1)

if __name__ == "__main__":
    main()