import streamlit as st
import networkx as nx
from pyvis.network import Network
import json
import os

# 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)
        
        # Function to create a NetworkX graph from data
        def create_graph(data):
            G = nx.DiGraph()
            for node in data.get("nodes", []):
                G.add_node(node["id"], label=node.get("label", node["id"]))
            for edge in data.get("edges", []):
                G.add_edge(edge["source"], edge["target"], label=edge.get("label", ""))
            return G
        
        # Generate the graph
        graph = create_graph(graph_data)
        
        # 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)
        
        # Save and display the graph
        output_file = "graph.html"
        pyvis_graph.show(output_file)
        
        # Display the graph in Streamlit
        st.components.v1.html(open(output_file, "r").read(), height=600, width=800)
        
        # Clean up the temporary file
        if os.path.exists(output_file):
            os.remove(output_file)
    except Exception as e:
        st.error(f"Error processing the file: {e}")
else:
    st.info("Please upload a JSON file to display the graph.")