import streamlit as st
import networkx as nx
from pyvis.network import Network
import json
from streamlit.components.v1 import html

# Streamlit app layout
st.title("Interactive Graph Visualization")
st.write("Upload a JSON file containing nodes and edges to display the graph.")

# File upload
uploaded_file = st.file_uploader("Upload JSON file", type=["json"])

if uploaded_file is not None:
    try:
        # Load JSON data
        graph_data = json.load(uploaded_file)
        
        # Validate JSON structure
        if "nodes" not in graph_data or "edges" not in graph_data:
            raise ValueError("The JSON file must contain 'nodes' and 'edges' keys.")
        
        # Function to create a NetworkX graph from data
        def create_graph(data):
            G = nx.DiGraph()
            for node in data["nodes"]:
                G.add_node(node["id"], label=node.get("label", node["id"]))
            for edge in data["edges"]:
                G.add_edge(edge["source"], edge["target"], label=edge.get("label", ""))
            return G
        
        # Generate the graph
        graph = create_graph(graph_data)
        
        # Check if the graph has any nodes or edges
        if graph.number_of_nodes() == 0:
            raise ValueError("The graph has no nodes.")
        if graph.number_of_edges() == 0:
            raise ValueError("The graph has no edges.")
        
        # Function to create a pyvis network from NetworkX graph
        def create_pyvis_graph(G):
            net = Network(height="600px", width="100%", directed=True)
            for node, data in G.nodes(data=True):
                net.add_node(node, label=data.get("label", node))
            for source, target, data in G.edges(data=True):
                net.add_edge(source, target, title=data.get("label", ""))
            return net

        # Create the Pyvis graph
        pyvis_graph = create_pyvis_graph(graph)
        
        # Generate the HTML representation of the graph
        pyvis_graph_html = pyvis_graph.generate_html()

        # Display the graph in Streamlit
        html(pyvis_graph_html, height=600)

    except Exception as e:
        st.error(f"Error processing the file: {e}")
else:
    st.info("Please upload a JSON file to display the graph.")