####################################################
# DVrouter.py
# Name: Yuwei Wan, Haoran Xu
# JHED ID: ywan10, hxu64
#####################################################

import sys
from collections import defaultdict
from router import Router
from packet import Packet
from json import dumps, loads


class DVrouter(Router):
    """Distance vector routing protocol implementation."""

    def __init__(self, addr, heartbeatTime):
        """TODO: add your own class fields and initialization code here"""
        Router.__init__(self, addr)  # initialize superclass - don't remove
        self.heartbeatTime = heartbeatTime
        self.last_time = 0
        # Hints: initialize local state
        self.addr = addr
        self.forwardingTable = {addr: {
            "interface": None,
            'nextHop': None,
            'cost': 0
        }}
        self.back = {}
        self.INF = 16

    def broadcast(self):
        for t_addr, t_distanceVector in self.forwardingTable.iteritems():
            table = loads(dumps(self.forwardingTable))
            for k, v in table.iteritems():
                if k == v['nextHop']:
                    table[k] = {
                        'cost': self.INF,
                        'nextHop': None,
                        'interface': None
                    }
            packet = Packet(Packet.ROUTING, self.addr, t_addr,
                            dumps(table))
            self.send(t_distanceVector["interface"], packet)

    def handlePacket(self, port, packet):
        """TODO: process incoming packet"""
        if packet.isTraceroute():
            # Hints: this is a normal data packet
            # if the forwarding table contains packet.dstAddr
            #   send packet based on forwarding table, e.g., self.send(port, packet)
            if packet.dstAddr in self.forwardingTable:
                self.send(self.forwardingTable[packet.dstAddr]["interface"], packet)
        else:
            # Hints: this is a routing packet generated by your routing protocol
            # if the received distance vector is different
            #   update the local copy of the distance vector
            #   update the distance vector of this router
            #   update the forwarding table
            #   broadcast the distance vector of this router to neighbors
            DV = loads(packet.content)
            for addr, distanceVector in DV.iteritems():
                if addr != self.addr:
                    if addr not in self.forwardingTable:
                        if self.forwardingTable[packet.srcAddr]["interface"] is not None:
                            self.forwardingTable[addr] = {
                                "cost": self.forwardingTable[packet.srcAddr]["cost"] + distanceVector["cost"],
                                "nextHop": packet.srcAddr,
                                "interface": self.forwardingTable[packet.srcAddr]["interface"]
                            }
                    else:
                        if (self.forwardingTable[packet.srcAddr]["cost"] + distanceVector["cost"] <
                            self.forwardingTable[addr]["cost"]) and \
                                self.forwardingTable[packet.srcAddr]["interface"] is not None:
                            self.forwardingTable[addr] = {
                                "cost": self.forwardingTable[packet.srcAddr]["cost"] + distanceVector["cost"],
                                "nextHop": packet.srcAddr,
                                "interface": self.forwardingTable[packet.srcAddr]["interface"]
                            }
                        elif self.forwardingTable[addr]['nextHop'] == packet.srcAddr:
                            self.forwardingTable[addr] = {
                                "cost": self.forwardingTable[packet.srcAddr]["cost"] + distanceVector["cost"],
                                "nextHop": packet.srcAddr,
                                "interface": self.forwardingTable[packet.srcAddr]["interface"]
                            }
            if dumps(self.back) != dumps(self.forwardingTable):
                self.back = self.forwardingTable
                self.broadcast()

    def handleNewLink(self, port, endpoint, cost):
        """TODO: handle new link"""
        # update the distance vector of this router
        # update the forwarding table
        # broadcast the distance vector of this router to neighbors
        if endpoint not in self.forwardingTable or cost < self.forwardingTable[endpoint]["cost"]:
            self.forwardingTable[endpoint] = {"cost": cost, "nextHop": endpoint, "interface": port}
        self.back = self.forwardingTable
        self.broadcast()

    def handleRemoveLink(self, port):
        """TODO: handle removed link"""
        # update the distance vector of this router
        # update the forwarding table
        # broadcast the distance vector of this router to neighbors
        for addr, distanceVector in self.forwardingTable.iteritems():
            if distanceVector["interface"] == port:
                self.forwardingTable[addr] = {"cost": self.INF, "nextHop": None, "interface": None}
        self.back = self.forwardingTable
        self.broadcast()

    def handleTime(self, timeMillisecs):
        """TODO: handle current time"""
        if timeMillisecs - self.last_time >= self.heartbeatTime:
            self.last_time = timeMillisecs
            # broadcast the distance vector of this router to neighbors
            self.back = self.forwardingTable
            for addr, distanceVector in self.forwardingTable.iteritems():
                packet = Packet(Packet.ROUTING, self.addr, addr, dumps(self.forwardingTable))
                self.send(distanceVector["interface"], packet)

    def debugString(self):
        """TODO: generate a string for debugging in network visualizer"""
        return "" + str(self.forwardingTable)
