#!/usr/bin/env python3
"""
Test Neo4j graph database connectivity and functionality
"""

import asyncio
from neo4j import AsyncGraphDatabase

# Configuration
NEO4J_URL = "bolt://localhost:7687"
NEO4J_USER = "neo4j"
NEO4J_PASSWORD = "langmem_neo4j_password"

async def test_neo4j_connection():
    """Test Neo4j connection and basic operations"""
    print("🧪 Testing Neo4j Graph Database")
    print("=" * 50)
    
    try:
        # Create driver
        driver = AsyncGraphDatabase.driver(NEO4J_URL, auth=(NEO4J_USER, NEO4J_PASSWORD))
        
        # Test connection
        async with driver.session() as session:
            print("1. Testing connection...")
            result = await session.run("RETURN 1 as test")
            record = await result.single()
            print(f"   ✅ Connection successful: {record['test']}")
            
            # Clear existing data
            print("\n2. Clearing existing data...")
            await session.run("MATCH (n) DETACH DELETE n")
            print("   ✅ Database cleared")
            
            # Create test memory nodes with relationships
            print("\n3. Creating test memory nodes...")
            
            # Create memory nodes
            memories = [
                {
                    "id": "mem1",
                    "content": "Claude Code is an AI-powered CLI tool",
                    "category": "tools",
                    "tags": ["claude", "ai", "cli"]
                },
                {
                    "id": "mem2", 
                    "content": "FastAPI is a modern web framework",
                    "category": "frameworks",
                    "tags": ["fastapi", "python", "web"]
                },
                {
                    "id": "mem3",
                    "content": "Docker provides containerization",
                    "category": "devops",
                    "tags": ["docker", "containers"]
                }
            ]
            
            for memory in memories:
                await session.run("""
                    CREATE (m:Memory {
                        id: $id,
                        content: $content,
                        category: $category,
                        tags: $tags,
                        created_at: datetime()
                    })
                """, **memory)
            
            print(f"   ✅ Created {len(memories)} memory nodes")
            
            # Create category nodes and relationships
            print("\n4. Creating category relationships...")
            await session.run("""
                MATCH (m:Memory)
                MERGE (c:Category {name: m.category})
                MERGE (m)-[:BELONGS_TO]->(c)
            """)
            
            # Create tag nodes and relationships
            await session.run("""
                MATCH (m:Memory)
                UNWIND m.tags as tag
                MERGE (t:Tag {name: tag})
                MERGE (m)-[:HAS_TAG]->(t)
            """)
            
            print("   ✅ Created category and tag relationships")
            
            # Create similarity relationships (example)
            print("\n5. Creating similarity relationships...")
            await session.run("""
                MATCH (m1:Memory), (m2:Memory)
                WHERE m1.id = 'mem1' AND m2.id = 'mem2'
                CREATE (m1)-[:SIMILAR_TO {score: 0.65, reason: 'both are development tools'}]->(m2)
            """)
            
            await session.run("""
                MATCH (m1:Memory), (m2:Memory)
                WHERE m1.id = 'mem2' AND m2.id = 'mem3'
                CREATE (m1)-[:RELATED_TO {score: 0.45, reason: 'both used in modern development'}]->(m2)
            """)
            
            print("   ✅ Created similarity relationships")
            
            # Test queries
            print("\n6. Testing graph queries...")
            
            # Count nodes
            result = await session.run("MATCH (n) RETURN labels(n) as label, count(n) as count")
            print("   Node counts:")
            async for record in result:
                print(f"     {record['label']}: {record['count']}")
            
            # Find memories by category
            result = await session.run("""
                MATCH (m:Memory)-[:BELONGS_TO]->(c:Category {name: 'tools'})
                RETURN m.content as content
            """)
            print("\n   Memories in 'tools' category:")
            async for record in result:
                print(f"     - {record['content']}")
            
            # Find similar memories
            result = await session.run("""
                MATCH (m1:Memory)-[r:SIMILAR_TO]->(m2:Memory)
                RETURN m1.content as memory1, m2.content as memory2, r.score as score
            """)
            print("\n   Similar memories:")
            async for record in result:
                print(f"     - {record['memory1'][:30]}... → {record['memory2'][:30]}... (score: {record['score']})")
            
            # Find memories with common tags
            result = await session.run("""
                MATCH (m1:Memory)-[:HAS_TAG]->(t:Tag)<-[:HAS_TAG]-(m2:Memory)
                WHERE m1.id < m2.id
                RETURN m1.content as memory1, m2.content as memory2, t.name as common_tag
            """)
            print("\n   Memories with common tags:")
            async for record in result:
                print(f"     - {record['memory1'][:30]}... & {record['memory2'][:30]}... (tag: {record['common_tag']})")
            
            print("\n7. Testing graph traversal...")
            # Complex traversal query
            result = await session.run("""
                MATCH path = (m:Memory)-[:HAS_TAG]->(t:Tag)<-[:HAS_TAG]-(related:Memory)
                WHERE m.id = 'mem1' AND m <> related
                RETURN related.content as related_content, t.name as via_tag
            """)
            print("   Memories related to 'mem1' via tags:")
            async for record in result:
                print(f"     - {record['related_content'][:40]}... (via tag: {record['via_tag']})")
                
        await driver.close()
        
        print("\n" + "=" * 50)
        print("🎉 Neo4j Graph Database Test Complete!")
        print("✅ All tests passed successfully")
        print(f"📊 Graph database is working with relationships and traversals")
        print(f"🌐 Neo4j Browser: http://localhost:7474")
        print(f"   Username: {NEO4J_USER}")
        print(f"   Password: {NEO4J_PASSWORD}")
        
    except Exception as e:
        print(f"❌ Neo4j test failed: {e}")
        return False
    
    return True

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