#!/usr/bin/python -u
# The -u makes output unbuffered, so it will show up immediately
import json
import select
import socket
import sys
import time
import numpy as np
import threading

# BPDU format:
#
# {"source": <bridge ID>, "dest": <ID or 'ffff'>, "type": <bpdu or data>, "message": <message data>}
# {"source":"02a1", "dest":"ffff", "type": "bpdu", "message":{"id":"02a1", "root":"02a1", "cost":0}}

# Bridge class for forwarding packets, broadcasting bpdus, and closing ports to build a spanning tree
class Bridge:

    def __init__(self, bridge_id, LANs):
        # only needs to be in the initialization of the bridge
        # creates sockets and connects to them
        def build_socket(lan):
            s = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
            padded_lan = '\0' + lan
            while len(padded_lan) < 108:
                padded_lan += '\0'
            s.connect(padded_lan)
            return s

        # TODO start using this
        self.last_received = {}

        self.forward_table = {}
        self.bridge_id = bridge_id
        self.root_id = bridge_id
        self.cost_to_root = 0
        self.lans = {sock.fileno():sock for sock in map(build_socket,LANs)}
        self.enabled = {port : True for port in self.lans}
        print("Bridge " + id + " starting up\n")

    def broadcast(self, msg, port_to_exclude=None):
        for port, sock in self.lans.iteritems():
            if self.enabled[port] and port is not port_to_exclude:
                sock.send(msg)

    def broadcast_task(self):
        while True:
            bpdu = Bridge.BPDU(self.bridge_id, self.root_id, self.cost_to_root)
            print("Broadcasting: {}/{}".format(self.root_id, self.cost_to_root))
            self.broadcast(str(bpdu))
            # TODO it will be better to use something like a timer class with the expiration stuff
            time.sleep(0.5)

    # determine whether a message needs to be forwarded or ignored
    def handle_data(self, port, data):
        if not self.enabled[port]:
            return
        source = data.get('source')
        dest = data.get('dest')
        self.forward_table[source] = port
        if dest not in self.forward_table:
            print("Broadcasting message to all but {}".format(port))
            self.broadcast(json.dumps(data), port)
        else:
            dest_port = self.forward_table[dest]
            if dest_port is not port:
                print("Forwarding message to {}".format(dest_port))
                self.lans[dest_port].send(json.dumps(data))
            # else ignore it

    def update_root(self, data):
        print("Updating root to {}/{}".format(data['message']['root'], int(data['message']['cost']) + 1))
        self.root_id = data['message']['root']
        self.cost_to_root = int(data['message']['cost']) + 1
        # Clear forwarding table?
        self.forward_table = {}
        for p in self.enabled:
            self.enabled[p] = True

    # manage the bpdu message to build this bridges place in the spanning tree
    def handle_bpdu(self, port, data):
        self.last_received[data.get("source")] = time.time()
        b_root_val = int(data['message']['root'],16)
        b_bridge_val = int(data['message']['id'],16)
        b_cost = int(data['message']['cost'])
        root_val = int(self.root_id,16)

        if (b_root_val >= root_val) and (b_cost + 1 > self.cost_to_root) and (b_bridge_val > int(self.bridge_id, 16)):
            pass
        else:
            self.update_root(data)

        # Disable the port connecting these bridges if they already share the same cost and root
        # elif (b_root_val == root_val and b_cost == self.cost_to_root):# and b_bridge_val < root_val):
        if self.enabled[port]:
            if (b_root_val == root_val and b_cost +1 <= self.cost_to_root and int(self.bridge_id, 16) > b_bridge_val):
                print('Disabling port {}'.format(port))
                self.enabled[port] = False

    # receive messages and then delegate as data or bpdu
    def receive(self):
        # Calls select with all the sockets; change the timeout value (1)
        ready, ignore, ignore2 = select.select(self.lans.values(), [], [], 1)
        # Reads from each fo the ready sockets
        for sock in ready:
            data = sock.recv(1500)
            data = json.loads(data)
            print(data)
            if data.get("type") == "data":
                self.handle_data(sock.fileno(), data)
            else:
                self.handle_bpdu(sock.fileno(), data)

    # start the broadcast thread and the receiving loop
    def start(self):
        try:
            threading.Thread(target=self.broadcast_task).start()
            while True:
                self.receive()
        except:
            print(sys.exc_info())

    # A BPDU class to keep the information about this bridge message
    class BPDU:
        def __init__(self, bridge_id, root_id, cost_to_root, dest='ffff'):
            self.bridge_id = bridge_id
            self.dest = dest
            self.root_id = root_id
            self.cost_to_root = cost_to_root
        def __str__(self):
            msg = {
                "source": self.bridge_id,
                "dest": self.dest,
                "type": "bpdu",
                "message": {
                    "id": self.bridge_id,
                    "root": self.root_id,
                    "cost": self.cost_to_root
                }
            }
            return json.dumps(msg)

if __name__ == "__main__":
    id = sys.argv[1]
    LAN = sys.argv[2:]
    # TODO check if this is correct that we want only distinct lan values
    bridge = Bridge(id, np.unique(LAN))
    bridge.start()
