#!/usr/bin/env python3
"""
Comprehensive test of mem0 with database storage verification
"""

import sys
import psycopg2
import json
from config import load_config, get_mem0_config
from mem0 import Memory

def verify_database_storage():
    """Test mem0 functionality and verify data storage in database"""
    print("=" * 70)
    print("MEM0 DATABASE STORAGE VERIFICATION TEST")
    print("=" * 70)
    
    connection_string = "postgresql://supabase_admin:CzkaYmRvc26Y@localhost:5435/postgres"
    
    try:
        # Initialize mem0
        system_config = load_config()
        config = get_mem0_config(system_config, "ollama")
        print(f"🔧 Configuration: {json.dumps(config, indent=2)}")
        
        print("\n🚀 Initializing mem0...")
        memory = Memory.from_config(config)
        print("✅ mem0 initialized successfully")
        
        # Connect to database for verification
        print("\n🔌 Connecting to Supabase database...")
        db_conn = psycopg2.connect(connection_string)
        db_cursor = db_conn.cursor()
        print("✅ Database connection established")
        
        # Check initial state
        print("\n📊 Initial database state:")
        db_cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'vecs';")
        tables = db_cursor.fetchall()
        print(f"   Tables in vecs schema: {[t[0] for t in tables]}")
        
        test_user = "database_test_user"
        
        # Test 1: Add memories and verify storage
        print(f"\n📝 TEST 1: Adding memories for user '{test_user}'...")
        test_memories = [
            "I am passionate about artificial intelligence and machine learning",
            "My favorite programming language is Python for data science",
            "I enjoy working with vector databases like Supabase",
            "Local LLM models like Ollama are very impressive",
            "Graph databases such as Neo4j are great for relationships"
        ]
        
        for i, content in enumerate(test_memories):
            print(f"   Adding memory {i+1}: {content[:50]}...")
            result = memory.add(content, user_id=test_user)
            print(f"      Result: {result}")
            
            # Verify in database immediately
            collection_name = config['vector_store']['config']['collection_name']
            db_cursor.execute(f"""
                SELECT id, metadata, vec 
                FROM vecs."{collection_name}" 
                WHERE metadata->>'user_id' = %s
            """, (test_user,))
            records = db_cursor.fetchall()
            print(f"      Database records for {test_user}: {len(records)}")
        
        # Test 2: Verify complete database state
        print(f"\n📊 TEST 2: Complete database verification...")
        db_cursor.execute(f'SELECT COUNT(*) FROM vecs."{collection_name}";')
        total_count = db_cursor.fetchone()[0]
        print(f"   Total records in collection: {total_count}")
        
        db_cursor.execute(f"""
            SELECT id, metadata->>'user_id' as user_id, metadata, 
                   CASE 
                     WHEN vec IS NOT NULL THEN 'Vector stored' 
                     ELSE 'No vector'
                   END as vector_status
            FROM vecs."{collection_name}" 
            WHERE metadata->>'user_id' = %s
        """, (test_user,))
        
        user_records = db_cursor.fetchall()
        print(f"   Records for {test_user}: {len(user_records)}")
        for record in user_records:
            metadata = record[2] if record[2] else {}
            content = metadata.get('content', 'N/A') if isinstance(metadata, dict) else 'N/A'
            print(f"      ID: {record[0][:8]}... | User: {record[1]} | Vector: {record[3]} | Content: {str(content)[:40]}...")
        
        # Test 3: Search functionality
        print(f"\n🔍 TEST 3: Search functionality...")
        search_queries = [
            "artificial intelligence",
            "Python programming", 
            "vector database",
            "machine learning"
        ]
        
        for query in search_queries:
            print(f"   Search: '{query}'")
            results = memory.search(query, user_id=test_user)
            if 'results' in results and results['results']:
                for j, result in enumerate(results['results'][:2]):
                    score = result.get('score', 0)
                    memory_text = result.get('memory', 'N/A')
                    print(f"      {j+1}. Score: {score:.3f} | {memory_text[:50]}...")
            else:
                print("      No results found")
        
        # Test 4: Get all memories
        print(f"\n📋 TEST 4: Retrieve all memories...")
        all_memories = memory.get_all(user_id=test_user)
        if 'results' in all_memories:
            memories_list = all_memories['results']
            print(f"   Retrieved {len(memories_list)} memories via mem0")
            for i, mem in enumerate(memories_list[:3]):  # Show first 3
                print(f"      {i+1}. {mem.get('memory', 'N/A')[:50]}...")
        else:
            print(f"   Unexpected format: {all_memories}")
        
        # Test 5: Cross-verify with direct database query
        print(f"\n🔄 TEST 5: Cross-verification with database...")
        db_cursor.execute(f"""
            SELECT metadata
            FROM vecs."{collection_name}" 
            WHERE metadata->>'user_id' = %s
            ORDER BY (metadata->>'created_at')::timestamp
        """, (test_user,))
        
        db_records = db_cursor.fetchall()
        print(f"   Direct database query found {len(db_records)} records")
        
        # Verify data consistency
        mem0_count = len(all_memories.get('results', []))
        db_count = len(db_records)
        if mem0_count == db_count:
            print(f"   ✅ Data consistency verified: {mem0_count} records match")
        else:
            print(f"   ⚠️  Data mismatch: mem0={mem0_count}, database={db_count}")
        
        # Test 6: Test different users
        print(f"\n👥 TEST 6: Multi-user testing...")
        other_user = "other_test_user"
        memory.add("This is a memory for a different user", user_id=other_user)
        
        # Verify user isolation
        user1_memories = memory.get_all(user_id=test_user)
        user2_memories = memory.get_all(user_id=other_user)
        
        print(f"   User '{test_user}': {len(user1_memories.get('results', []))} memories")
        print(f"   User '{other_user}': {len(user2_memories.get('results', []))} memories")
        
        # Test 7: Vector storage verification
        print(f"\n🧮 TEST 7: Vector storage verification...")
        db_cursor.execute(f"""
            SELECT id, 
                   CASE WHEN vec IS NOT NULL THEN 'Vector present' ELSE 'No vector' END as vec_status,
                   substring(vec::text, 1, 100) as vec_preview
            FROM vecs."{collection_name}" 
            LIMIT 3
        """)
        
        vectors = db_cursor.fetchall()
        print("   Vector storage details:")
        for vec in vectors:
            print(f"      ID: {vec[0][:8]}... | Status: {vec[1]} | Preview: {vec[2][:60]}...")
        
        # Cleanup
        print(f"\n🧹 CLEANUP: Removing test data...")
        try:
            # Delete test users' memories
            for mem in user1_memories.get('results', []):
                memory.delete(mem['id'])
            for mem in user2_memories.get('results', []):
                memory.delete(mem['id'])
            print("   ✅ Test data cleaned up")
        except Exception as e:
            print(f"   ⚠️  Cleanup warning: {e}")
        
        # Final verification
        db_cursor.execute(f'SELECT COUNT(*) FROM vecs."{collection_name}";')
        final_count = db_cursor.fetchone()[0]
        print(f"   Final record count: {final_count}")
        
        db_cursor.close()
        db_conn.close()
        
        print("\n" + "=" * 70)
        print("🎉 ALL DATABASE STORAGE TESTS PASSED!")
        print("✅ Data is being properly stored in Supabase")
        print("✅ Vector embeddings are correctly stored")
        print("✅ User isolation working")
        print("✅ Search functionality operational")
        print("=" * 70)
        return True
        
    except Exception as e:
        print(f"\n❌ Test failed: {str(e)}")
        print(f"Error type: {type(e).__name__}")
        import traceback
        traceback.print_exc()
        return False

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