#!/usr/bin/env python3
"""
Test LangMem API integration with Neo4j graph relationships
"""

import asyncio
import httpx
import json

# Configuration
API_BASE_URL = "http://localhost:8765"
API_KEY = "langmem_api_key_2025"

async def test_langmem_with_neo4j():
    """Test LangMem API with Neo4j graph relationships"""
    print("🧪 Testing LangMem API + Neo4j Integration")
    print("=" * 50)
    
    headers = {"Authorization": f"Bearer {API_KEY}"}
    
    async with httpx.AsyncClient() as client:
        # Store a memory with graph relationships
        print("\n1. Storing memory with graph relationships...")
        
        memory_with_relationships = {
            "content": "LangMem is a long-term memory system for LLM projects that combines vector search with graph relationships",
            "user_id": "graph_test_user",
            "session_id": "graph_test_session",
            "metadata": {
                "category": "ai_systems",
                "subcategory": "memory_systems",
                "importance": "high",
                "tags": ["langmem", "llm", "vector", "graph", "memory"]
            },
            "relationships": [
                {
                    "entity_name": "Vector Search",
                    "entity_type": "Technology",
                    "relationship": "USES",
                    "confidence": 0.95,
                    "properties": {
                        "implementation": "pgvector",
                        "embedding_model": "nomic-embed-text"
                    }
                },
                {
                    "entity_name": "Graph Database",
                    "entity_type": "Technology", 
                    "relationship": "USES",
                    "confidence": 0.90,
                    "properties": {
                        "implementation": "neo4j",
                        "query_language": "cypher"
                    }
                },
                {
                    "entity_name": "LLM Projects",
                    "entity_type": "Domain",
                    "relationship": "SERVES",
                    "confidence": 0.98,
                    "properties": {
                        "purpose": "long_term_memory",
                        "integration": "mcp_server"
                    }
                }
            ]
        }
        
        try:
            response = await client.post(
                f"{API_BASE_URL}/v1/memories/store",
                json=memory_with_relationships,
                headers=headers,
                timeout=30.0
            )
            
            if response.status_code == 200:
                data = response.json()
                memory_id = data["id"]
                print(f"   ✅ Memory stored with ID: {memory_id}")
                print(f"   ✅ Graph relationships will be processed in background")
                
                # Wait a moment for background task
                await asyncio.sleep(2)
                
                # Search for the memory with graph relationships
                print("\n2. Searching memory with graph relationships...")
                
                search_response = await client.post(
                    f"{API_BASE_URL}/v1/memories/search",
                    json={
                        "query": "memory system for AI projects",
                        "user_id": "graph_test_user",
                        "limit": 5,
                        "threshold": 0.5,
                        "include_graph": True
                    },
                    headers=headers,
                    timeout=30.0
                )
                
                if search_response.status_code == 200:
                    search_data = search_response.json()
                    print(f"   ✅ Found {search_data['total_count']} memories")
                    
                    for memory in search_data['memories']:
                        print(f"     - Content: {memory['content'][:60]}...")
                        print(f"       Similarity: {memory['similarity']:.3f}")
                        if 'relationships' in memory:
                            print(f"       Relationships: {len(memory['relationships'])}")
                            for rel in memory['relationships']:
                                print(f"         → {rel['relationship']} {rel['entity_name']} ({rel['confidence']})")
                        else:
                            print("       Relationships: Not included")
                else:
                    print(f"   ❌ Search failed: {search_response.status_code}")
                    
            else:
                print(f"   ❌ Memory storage failed: {response.status_code}")
                print(f"   Response: {response.text}")
                
        except Exception as e:
            print(f"   ❌ Error: {e}")
        
        # Test retrieval for conversation
        print("\n3. Testing memory retrieval with graph context...")
        
        try:
            retrieval_response = await client.post(
                f"{API_BASE_URL}/v1/memories/retrieve",
                json={
                    "messages": [
                        {"role": "user", "content": "Tell me about memory systems"},
                        {"role": "assistant", "content": "I can help with memory systems. What specific aspect?"},
                        {"role": "user", "content": "How do vector databases work with graph relationships?"}
                    ],
                    "user_id": "graph_test_user",
                    "session_id": "graph_test_session"
                },
                headers=headers,
                timeout=30.0
            )
            
            if retrieval_response.status_code == 200:
                retrieval_data = retrieval_response.json()
                print(f"   ✅ Retrieved {retrieval_data['total_count']} relevant memories")
                
                for memory in retrieval_data['memories']:
                    print(f"     - {memory['content'][:50]}... (similarity: {memory['similarity']:.3f})")
                    if 'relationships' in memory:
                        print(f"       Graph relationships: {len(memory['relationships'])}")
            else:
                print(f"   ❌ Retrieval failed: {retrieval_response.status_code}")
                
        except Exception as e:
            print(f"   ❌ Retrieval error: {e}")
        
        print("\n" + "=" * 50)
        print("🎉 LangMem + Neo4j Integration Test Complete!")
        print("✅ Vector search and graph relationships working together")
        print("🌐 Check Neo4j Browser: http://localhost:7474")
        print("   Look for Memory, Entity nodes and their relationships")

if __name__ == "__main__":
    asyncio.run(test_langmem_with_neo4j())