#!/usr/bin/python3

import sys, socket, select

seqNum = 0
topoview = {} # (u,v) -> cost
topochanged = False
nodes = [] # [nodeIDs]
recvdAds = [] # (aID, seq)
myNeighbor = [] # (nID, cost)
forwardTable = {} # dest -> (next, totalcost)

# check arguments
if (len(sys.argv) != 4):
    print("Usage: virtualrouter <nfe-ip> <nfe-port> <virtual-router-id>", file=sys.stderr)
    sys.exit(0)

if (not sys.argv[2].isnumeric() or not sys.argv[3].isnumeric()):
    print("<nfe-port> or <virtual-router-id> must be a number", file=sys.stderr)
    sys.exit(0)

#extract info from args
serverName = sys.argv[1]
serverPort = int(sys.argv[2])
routerId = int(sys.argv[3])
recvdAds.append((routerId,seqNum))

# create forwarding table output file
FwdTblOutFile = open("forwarding_"+str(routerId)+".out","w")

# make new socket
addr = (serverName, serverPort)
try:
    nfeSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
except socket.error:
    print("Cannot create socket.", file=sys.stderr)
    sys.exit(1)

def isAdjacent(u, v):
    if u == v:
        return False
    uvpair = (min(u,v),max(u,v))
    if uvpair in topoview:
        return True
    return False

def linkCost(u,v):
    if u == v:
        return 0
    uvpair = (min(u,v),max(u,v))
    if uvpair in topoview:
        return topoview[uvpair]
    return 2147483647

def Dijkstra():
    V = nodes[:] #take a copy of nodes
    V.remove(routerId) # V is complement of N' # N' = {u}
    D = {} # dest -> (next, totalcost)
    V.sort()
    #Initialization: 
    for v in nodes:
        if v != routerId:
            D[str(v)]=(v,linkCost(routerId,v))
    #Loop:
    while V:
        current = 0
        minD = 2147483647
        for w in V: # for w not in N' s.t Dw is minimum
            Dw = D[str(w)][1]
            if minD > Dw: #Dw is minimum
                minD = Dw
                current = w
        V.remove(current) # add w to N'
        for v in V: # for all v not in N'
            if (isAdjacent(v,current)): #adjacent to w
                Dv = D[str(v)][1]
                Dwcwv = D[str(current)][1] + linkCost(current,v)
                if Dwcwv < Dv: # min(D(v),D(w)+c(w,v))
                    D[str(v)] = (D[str(current)][0],Dwcwv)
    return D

def printOutForwardTable(ft):
    FwdTblOutFile.write("FORWARDING\n")
    print("FORWARDING", file=sys.stderr)
    for entry in ft.items():
        if entry[0] != str(routerId):
            linestr = entry[0]+":"+str(entry[1][0])+","+str(entry[1][1])
            FwdTblOutFile.write(linestr+"\n")
            print(linestr, file=sys.stderr)
    FwdTblOutFile.write("\n")

#prepare and send init
tosend = (1).to_bytes(4, byteorder='big')
tosend += routerId.to_bytes(4, byteorder='big')
nfeSocket.sendto(tosend, addr)

while 1: 
    readlist = [nfeSocket]
    readable, _, exc = select.select(readlist,[],[]) #no busy waiting

    for r in readable:
        data, addr = r.recvfrom(4096)
        # print(len(data))
        packettype = int.from_bytes(data[0:4], byteorder='big')
        # print(packettype)
        if (packettype == 3): #LSA
            sendID = int.from_bytes(data[4:8], byteorder='big')
            outstringSID = "SID(" + str(sendID) +"),"
            destID = int.from_bytes(data[8:12], byteorder='big')
            outstringDID = "DID("+str(destID)+"),"
            advID = int.from_bytes(data[12:16], byteorder='big')
            outstringTail = "AID("+str(advID)+"),"
            seq = int.from_bytes(data[16:20], byteorder='big')
            outstringTail += "SEQ("+str(seq)+")"
            numLink = int.from_bytes(data[20:24], byteorder='big')
            if advID not in nodes:
                nodes.append(advID)
            for i in range(numLink): # extract link info
                nID = int.from_bytes(data[24+i*8:28+i*8], byteorder='big')
                cost = int.from_bytes(data[28+i*8:32+i*8], byteorder='big')
                outstringTail += ",NID("+str(nID)+"),LC("+str(cost)+")"
                link = (min(advID,nID),max(advID,nID))
                if nID not in nodes:
                    nodes.append(nID)
                if link not in topoview:
                    topochanged = True
                    topoview[link] = cost
            print("RECEIVED:" + outstringSID + outstringDID + outstringTail)
            if (advID,seq) in recvdAds: # drop this LSA packet
                print("DROPPING:" + outstringSID + outstringDID + outstringTail)
                continue
            else: # mark it as received, update forwarding, send to neighbors
                recvdAds.append((advID,seq))
                outstringSID = "SID(" + str(routerId) +"),"
                if topochanged: # run Dijkstra and print forward table
                    forwardTable = Dijkstra()
                    printOutForwardTable(forwardTable)
                    topochanged = False
                forwardHead = data[0:4] + routerId.to_bytes(4, byteorder='big') #msg type + sender ID
                forwardUnchanged = data[12:len(data)]
                for n in myNeighbor:
                    forwardDest = n[0].to_bytes(4, byteorder='big') # dest ID
                    if forwardDest != sendID:
                        outstringDID = "DID("+str(n[0])+"),"
                        print("FORWARDING:" + outstringSID + outstringDID + outstringTail)
                        nfeSocket.sendto(forwardHead+forwardDest+forwardUnchanged,addr)
        elif (packettype == 4): #init-reply
            numLink = data[4:8]
            numLink = int.from_bytes(numLink, byteorder='big')
            nodes.append(routerId)
            for i in range(numLink): # extract link info
                nID = int.from_bytes(data[8+i*8:12+i*8], byteorder='big')
                cost = int.from_bytes(data[12+i*8:16+i*8], byteorder='big')
                myNeighbor.append((nID,cost))
                link = (min(routerId,nID),max(routerId,nID))
                if link not in topoview:
                    topochanged = True
                    topoview[link] = cost
                    nodes.append(nID)
                # print(link, file=sys.stderr)
            # prepare and send LSAs to my neighbors
            LSAhead = (3).to_bytes(4, byteorder='big') #Message Type
            LSAhead += routerId.to_bytes(4, byteorder='big') #sender ID
            LSAunchanged = routerId.to_bytes(4, byteorder='big') #advID
            LSAunchanged += seqNum.to_bytes(4, byteorder='big') #sequence
            LSAunchanged += len(myNeighbor).to_bytes(4, byteorder='big') #numLinks
            outstringSID = "SID(" + str(routerId) + "),"
            outstringTail = "AID("+str(routerId)+"),SEQ("+str(seqNum)+")"
            for n in myNeighbor: # populate link info
                LSAunchanged += n[0].to_bytes(4, byteorder='big') #neighborID
                LSAunchanged += n[1].to_bytes(4, byteorder='big') #cost
                outstringTail += ",NID("+str(n[0])+"),LC("+str(n[1])+")"
            # print("LSA Prepared, now sending", file=sys.stderr)
            for n in myNeighbor: # send to each neighbor
                destinationID = n[0].to_bytes(4, byteorder='big') #destination
                outstringDID = "DID(" + str(n[0]) + "),"
                tosend = LSAhead + destinationID + LSAunchanged
                print("SENDING:"+outstringSID+outstringDID+outstringTail)
                nfeSocket.sendto(tosend,addr)
            if topochanged: # Run Dijkstra and print forwarding table
                forwardTable = Dijkstra()
                printOutForwardTable(forwardTable)
                topochanged = False
        elif (packettype == 5):
            FwdTblOutFile.close()
            sys.exit(0)

