# graph_builder.py
import networkx as nx
import json
import pickle
from typing import Dict, List, Any

class LegalProcessGraph:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.processes = {}
        
    def load_process_data(self, json_file_path: str):
        """Load process data from JSON file and build graph"""
        with open(json_file_path, 'r', encoding='utf-8') as file:
            process_data = json.load(file)
            
        process_name = process_data['process_name']
        self.processes[process_name] = process_data
        
        # Add nodes to graph
        for node in process_data['nodes']:
            self.graph.add_node(
                node['id'],
                **{
                    'title': node['title'],
                    'description': node['description'],
                    'type': node['type'],
                    'process': process_name,
                    'properties': node.get('properties', {})
                }
            )
            
        # Add edges to graph  
        for edge in process_data['edges']:
            self.graph.add_edge(
                edge['source'],
                edge['target'],
                relationship=edge['relationship'],
                condition=edge.get('condition')
            )
            
        # Add resource nodes
        for resource in process_data.get('resources', []):
            self.graph.add_node(
                resource['id'],
                **{
                    'title': resource['title'],
                    'type': resource['type'], 
                    'process': process_name,
                    'properties': {
                        'url': resource.get('url'),
                        'phone': resource.get('phone'),
                        'email': resource.get('email'),
                        'description': resource['description']
                    }
                }
            )
            
    def find_process_start(self, process_name: str) -> str:
        """Find the starting node for a given process"""
        for node_id, data in self.graph.nodes(data=True):
            if data.get('process') == process_name and data.get('type') == 'process_start':
                return node_id
        return None
        
    def get_next_steps(self, node_id: str) -> List[Dict]:
        """Get next steps from current node"""
        next_steps = []
        for successor in self.graph.successors(node_id):
            node_data = self.graph.nodes[successor]
            edge_data = self.graph.edges[node_id, successor]
            next_steps.append({
                'node_id': successor,
                'title': node_data['title'],
                'description': node_data['description'],
                'relationship': edge_data['relationship'],
                'condition': edge_data.get('condition'),
                'properties': node_data.get('properties', {})
            })
        return next_steps
        
    def get_required_resources(self, node_id: str) -> List[Dict]:
        """Get resources related to a specific node"""
        node_data = self.graph.nodes[node_id]
        process_name = node_data['process']
        
        # Find all resource nodes for this process
        resources = []
        for resource_id, data in self.graph.nodes(data=True):
            if data.get('process') == process_name and data.get('type') in ['government_portal', 'helpline', 'form']:
                resources.append({
                    'id': resource_id,
                    'title': data['title'],
                    'type': data['type'],
                    'properties': data.get('properties', {})
                })
        return resources
        
    def save_graph(self, file_path: str):
        """Save graph to pickle file"""
        with open(file_path, 'wb') as f:
            pickle.dump({
                'graph': self.graph,
                'processes': self.processes
            }, f)
            
    def load_graph(self, file_path: str):
        """Load graph from pickle file"""
        with open(file_path, 'rb') as f:
            data = pickle.load(f)
            self.graph = data['graph']
            self.processes = data['processes']

# Usage example
if __name__ == "__main__":
    # Build the graph
    legal_graph = LegalProcessGraph()
    legal_graph.load_process_data('legal_aid_process.json')
    legal_graph.save_graph('legal_processes.pkl')
    
    print(f"Graph built with {legal_graph.graph.number_of_nodes()} nodes and {legal_graph.graph.number_of_edges()} edges")