import json
from pathlib import Path
from typing import List, Dict, Set

class RLMTraversal:
    """
    Genesis Phase 4: Relationship Traversal.
    Enables deep graph exploration for semantic reasoning.
    """
    def __init__(self, rlm_bridge):
        self.rlm = rlm_bridge

    def get_related_nodes(self, node_name: str, relation_type: str = None, depth: int = 1) -> Set[str]:
        """Find related nodes up to N depth via PostgreSQL."""
        # Now using the pg_store direct connection if available
        pg = self.rlm.pg_store
        visited = {node_name}
        current_layer = {node_name}
        
        for _ in range(depth):
            next_layer = set()
            for node in current_layer:
                # In PG, we look for connections between episodes or entities
                # For this implementation, we query the search endpoint
                related = pg.get_connected_episodes(node, relation_type=relation_type)
                for rel in related:
                    target = rel.get("episode_id") or rel.get("entity_id")
                    if target and target not in visited:
                        next_layer.add(target)
                        visited.add(target)
            current_layer = next_layer
            if not current_layer: break
            
        return visited - {node_name}

    def get_context_cluster(self, node_name: str) -> Dict[str, List[str]]:
        """Retrieve a cluster of related nodes and their observations."""
        related = self.get_related_nodes(node_name, depth=2)
        graph = self.rlm.load_graph()
        
        cluster = {}
        for entity in graph.get("entities", []):
            if entity["name"] in related or entity["name"] == node_name:
                cluster[entity["name"]] = entity["observations"]
                
        return cluster