import sqlite3
import networkx as nx
from pyvis.network import Network
import matplotlib.pyplot as plt
from datetime import datetime

def retrieve_traceroute_data():
    conn = sqlite3.connect('traceroute.db')
    cursor = conn.cursor()
    cursor.execute('''
        SELECT l.source_ip, l.destination_ip, ll.latency_ms, ll.timestamp
        FROM Latency ll
        JOIN links l ON ll.link_id = l.id
        ORDER BY ll.timestamp
    ''')
    rows = cursor.fetchall()
    conn.close()
    return rows

def build_graph(data):
    G = nx.DiGraph()  # Create a directed graph
    for row in data:
        source_ip = row[0]
        destination_ip = row[1]
        latency = row[2]
        timestamp = row[3]

        # Add nodes and edges to the graph
        G.add_node(source_ip)
        G.add_node(destination_ip)
        G.add_edge(source_ip, destination_ip, latency=latency, timestamp=timestamp)

    return G

#def visualize_graph_pyvis(G, output_file='network.html'):
#    net = Network(height='750px', width='100%', directed=True)
#    net.from_nx(G)
#
#    for edge in G.edges(data=True):
#        src, dst, data = edge
#        latency = data['latency']
#        timestamp = data['timestamp']
#        net.add_edge(src, dst, title=f'Latency: {latency} ms<br>Timestamp: {timestamp}', value=latency)
#
#    net.show(output_file)

def main():
    # Retrieve data from the database
    traceroute_data = retrieve_traceroute_data()

    # Build the network graph
    graph = build_graph(traceroute_data)

    return graph

if __name__ == '__main__':
    graph = main()

    nx.draw_planar(graph, with_labels=True)
    plt.savefig("path.png")


    # Visualize the graph using pyvis
    #visualize_graph_pyvis(graph, 'traceroute_network.html')

#exit(0)