#!/usr/bin/python
# -*- coding: utf-8 -*-

import requests
from requests.auth import HTTPBasicAuth
import json
import networkx as nx

IP = '127.0.0.1'
PORT = 8080
USER = ''
PASS = ''
TOPO_UUID = 'ols-topo'

TOPO_LIST_URL = 'http://{:s}:{:d}/restconf/data/tapi-common:context/tapi-topology:topology-context/topology={:s}/'

def retrieveTopology(ip, port, user, password, topo_uuid):
    print ("Reading network-topology")
    topologies = []
    topo_list_url = TOPO_LIST_URL.format(ip, port, topo_uuid)
    response = requests.get(topo_list_url, auth=HTTPBasicAuth(user, password))
    topologies.append(response.json())
    print ("Retrieved Topology: " + json.dumps(response.json(), indent=4))
    return topologies

def draw_topologies (topologies):
    nwk_graph = nx.DiGraph()

    for topo in topologies:
        topo_uuid = topo['uuid']
        for node in topo['node']:
            if node['owned-node-edge-point']:
                nwk_graph.add_node(node['uuid'])

        for link in topo['link']:
            node1 = link['node-edge-point'][0]['node-uuid']
            node2 = link['node-edge-point'][1]['node-uuid']
            nwk_graph.add_edge(node1, node2)

        # Using Matplotlib
        #import matplotlib.pyplot as plt
        #nx.draw(nwk_graph, pos=nx.spring_layout(nwk_graph, scale=500))
        #plt.show(block=False)
        #plt.savefig('{:s}.png'.format(TOPO_UUID), format='PNG')

        # Using Pydot and Graphviz
        from networkx.drawing.nx_pydot import write_dot, to_pydot
        write_dot(nwk_graph, '{:s}.dot'.format(topo_uuid))
        dot_graph = to_pydot(nwk_graph)
        with open('{:s}.png'.format(topo_uuid), 'wb') as f:
            f.write(dot_graph.create(format='png'))

if __name__ == "__main__":
    topologies = retrieveTopology(IP, PORT, USER, PASS, TOPO_UUID)
    draw_topologies(topologies)