from neo4j import GraphDatabase
import time

uri = "bolt://localhost:7687"
driver = GraphDatabase.driver(uri, auth=("neo4j", "test"))

def add_event(tx, event_id, event_type, content, timestamp):
    tx.run(
        "MERGE (e:Event {id: $event_id, type: $event_type, content: $content, ts: $timestamp})",
        event_id=event_id, event_type=event_type, content=content, timestamp=timestamp
    )

def link_events(tx, from_id, to_id):
    tx.run(
        "MATCH (a:Event {id: $from_id}), (b:Event {id: $to_id}) "
        "MERGE (a)-[:RELATED_TO]->(b)",
        from_id=from_id, to_id=to_id,
    )

with driver.session() as session:
    start = time.time()
    for i in range(5000):
        session.write_transaction(add_event, f"evt_{i}", "summary", f"content_{i}", i)
    for i in range(1, 5000):
        session.write_transaction(link_events, f"evt_{i-1}", f"evt_{i}")
    latency = (time.time() - start) * 1000
    print(f"Neo4j insert + link 5k events: {latency:.2f} ms")

    query_start = time.time()
    res = session.run(
        "MATCH (e:Event)-[:RELATED_TO*1..3]->(other) WHERE e.id='evt_1000' "
        "RETURN other LIMIT 5"
    )
    print(f"3-hop query latency: {(time.time() - query_start) * 1000:.2f} ms")