import networkx as nx
import yaml

from networkx.drawing.nx_agraph import to_agraph

class Topology(object):
    def __init__(self):
        self.G = nx.MultiGraph()

    def render(self):
        '''
            Summary:
            Builds a network topology from the topology.yaml file into memory, stored in a networkx Graph object.
            Once the networkx Graph object has been built, it's drawn and saved as topology.png.

            Nodes: network devices
            Edges: links between network devices
            Edge attributes: dict of attributes assigned to a specific edge
        '''
        topology_data = yaml.load(open('topology.yaml'), Loader=yaml.SafeLoader)

        for topology_name, topology_vars in topology_data.items():
            for node in topology_vars['nodes']:
                self.G.add_node(node)
                self.G.nodes[node]['shape'] = 'box'

            for edge_number, edge_attributes in topology_vars['edges'].items():
                self.G.add_edge(edge_attributes['a_end'], edge_attributes['b_end'])
                nx.set_edge_attributes(
                    self.G, 
                    {
                        (edge_attributes['a_end'], 
                        edge_attributes['b_end'], 
                        edge_attributes['edge_index']): edge_attributes
                    }
                )

        for node1, node2, edge_attributes in self.G.edges(data=True):
            edge_attributes['label'] = f"""
            {edge_attributes['a_end']}
            {edge_attributes['a_end_ip']}, {edge_attributes['a_end_interface']}\n
            {edge_attributes['b_end']}
            {edge_attributes['b_end_ip']}, {edge_attributes['b_end_interface']}
            """

        vis = to_agraph(self.G)
        vis.layout('dot')
        vis_path = 'topology/topology.png'
        vis.draw(vis_path)

    def get_dns_records(self):
        '''
            Summary:
            Creates & outputs Ubuntu DNS entries for each edge IP in the networkx Graph object.
            Enabling operators to easily visualise the path traffic has taken across a network.

            dom@upsrv01:~/Python/T2DNS$ traceroute 10.0.0.3 --resolve -I
            traceroute to 10.0.0.3 (10.0.0.3), 64 hops max
              1   192.168.0.249 (UBUNTU_eth0-->GigabitEthernet8_LAB-RELAY)  0.864ms  0.688ms  0.720ms
              2   10.0.0.3 (LAB-RELAY_GigabitEthernet1-->GigabitEthernet8_LAB-897VA)  1.431ms  1.372ms  1.363ms
        '''
        records = {}

        print('Creating DNS records...')
        for node1, node2, edge_attributes in self.G.edges(data=True):
            records.update({edge_attributes['b_end_ip']: f"{node1}_{edge_attributes['a_end_interface']}-->{edge_attributes['b_end_interface']}_{node2}"})
            records.update({edge_attributes['a_end_ip']: f"{node2}_{edge_attributes['b_end_interface']}-->{edge_attributes['a_end_interface']}_{node1}"})

        for ip, record in records.items():
            print(f'{ip} {record}')