#!/usr/bin/python3
import sys
import struct
import wrapper
import threading
import time
from wrapper import recv_from_any_link, send_to_link, get_switch_mac, get_interface_name

BPDU_DEST_MAC = "01:80:c2:00:00:00"
def is_multicast(mac: str) -> bool:
    octet = int(mac.split(":")[0], 16)
    return (octet & 1) != 0 and mac.lower() != "ff:ff:ff:ff:ff:ff"

def create_packet(root_bridge_ID, sender_bridge_ID, sender_path):
    dest_mac = bytes.fromhex("0180C2000000")
    llc_header = struct.pack("3b", 0x42, 0x42, 0x03)
    bpdu_header = struct.pack("!i", 0)
    bpdu_config = struct.pack("!3i", root_bridge_ID, sender_bridge_ID, sender_path)
    llc_len = len(llc_header + bpdu_header + bpdu_config)
    return (
        dest_mac + 
        get_switch_mac() + 
        struct.pack("!H", llc_len) +
        llc_header + 
        bpdu_header +
        bpdu_config
    )
# Define BPDU structure
class BPDU:
    def __init__(self, switch_priority, port_vlan, interfaces):
        self.root_bridge_id = switch_priority
        self.own_bridge_id = switch_priority
        self.root_port = -1
        self.path_cost = 0
        self.port_vlan = port_vlan
        self.interfaces = interfaces
        self.port_states = {}

        for port in self.interfaces:
            if self.port_vlan.get(get_interface_name(port), -1) == "T":
                self.port_states[port] = "BLOCKING"
    

    def create_packet(self):
        return create_packet(self.own_bridge_id, self.own_bridge_id, 0)

    def receiving_bpdu(self, interface, data, length):
        root_bridge_id = struct.unpack("!H", data[0:2])[0]
        port_bridge_id = struct.unpack("!H", data[2:4])[0]
        path_cost = struct.unpack("!H", data[4:6])[0]
        if root_bridge_id < self.root_bridge_id:
            self.root_bridge_id = root_bridge_id
            self.path_cost = path_cost + 10
            self.path_cost = path_cost + 10
            self.root_port = interface

            if self.own_bridge_id == root_bridge_id:
                for port in self.interfaces:
                    if port != interface:
                        if self.port_vlan[get_interface_name(port)] == "T":
                            self.port_states[port] = "BLOCKING"
            if self.port_states[self.root_port] == "BLOCKING":
                self.port_states[self.root_port] = "LISTENING"
            
            new_data = create_packet(self.root_bridge_id, self.own_bridge_id, self.path_cost)
            for port in self.interfaces:
                if port != interface:
                    if self.port_vlan[get_interface_name(port)] == "T":
                        send_to_link(port, len(new_data), new_data)    

        elif root_bridge_id == self.root_bridge_id:
            if interface == self.root_port and path_cost + 10 < self.path_cost:
                self.path_cost = path_cost + 10
            elif interface != self.root_port:
                if path_cost > self.path_cost:
                    if self.port_states[interface] != "DESIGNATED":
                        self.port_stated[interface] = "DESIGNATED"

        elif self.own_bridge_id == port_bridge_id:
            self.port_states[interface] = "BLOCKING"

        if self.own_bridge_id == self.root_bridge_id:
            for port in self.port_states:
                if port != self.root_port and self.port_vlan[get_interface_name(port)] != "T":
                    self.port_states[port] = "DESIGNATED"

    def is_bpdu(self, interface, dest_mac, data, len) -> bool:
        if is_multicast(dest_mac):
            self.receiving_bpdu(interface, data, len)
            return True
        return False
    def not_blocked_interface(self):
        ports_unblocked = []
        for port in self.interfaces:
            if self.port_states[port] != "BLOCKING":
                ports_unblocked.append(port)
        return ports_unblocked   

def is_unicast(mac):
    least = int(mac.split(":")[0], 16)
    # if the least significant bit of the first byte is 0, then it is a unicast address
    return (least & 0x01) == 0

def send_vlan_packet(interface, data, length, port_vlan, vlan_id, vlan_in):
    vlan_out = port_vlan[get_interface_name(interface)]

    if vlan_out == "T":
        if vlan_id == -1:
            data = data[:12] + create_vlan_tag(vlan_in) + data[12:]
            length = len(data)
            send_to_link(interface, length, data)
        else:
            send_to_link(interface, length, data)
    elif int(vlan_out) == vlan_in:
        if vlan_id != -1:
            data = remove_vlan_tag(data)
            length = len(data)
        send_to_link(interface, length, data)             

def create_vlan_tag(vlan_id):
    tpid = 0x8200
    vlan_tci = vlan_id & 0x0FFF
    vlan_tag = struct.pack('!HH', tpid, vlan_tci)
    return vlan_tag

def remove_vlan_tag(data):
    return data[:12] + data[16:]    

def get_ports_data(switch_id):
    filename = f"configs/switch{switch_id}.cfg"
    port_vlan = {}
    switch_priority = 0

    with open(filename, "r") as f:
        lines = f.readlines()
        priority = int(lines[0].strip())
        for line in lines[1:]:
            parts = line.strip().split(" ")
            port = parts[0]
            vlan_type = parts[1]
            port_vlan[port] = vlan_type

    return priority, port_vlan


def parse_ethernet_header(data):
    # Unpack the header fields from the byte array
    #dest_mac, src_mac, ethertype = struct.unpack('!6s6sH', data[:14])
    dest_mac = data[0:6]
    src_mac = data[6:12]
    
    # Extract ethertype. Under 802.1Q, this may be the bytes from the VLAN TAG
    ether_type = (data[12] << 8) + data[13]

    vlan_id = -1
    # Check for VLAN tag (0x8100 in network byte order is b'\x81\x00')
    if ether_type == 0x8200:
        vlan_tci = int.from_bytes(data[14:16], byteorder='big')
        vlan_id = vlan_tci & 0x0FFF  # extract the 12-bit VLAN ID
        ether_type = (data[16] << 8) + data[17]

    return dest_mac, src_mac, ether_type, vlan_id

def send_bdpu_every_sec(bpdu):
    while True:
        if bpdu.own_bridge_id == bpdu.root_bridge_id:
            for port in bpdu.interfaces:
                if bpdu.port_vlan.get(get_interface_name(port), -1) == "T":
                    new_data = bpdu.create_packet()
                    send_to_link(port, len(new_data), new_data)
        # TODO Send BDPU every second if necessary
        time.sleep(1)

def main():
    # init returns the max interface number. Our interfaces
    # are 0, 1, 2, ..., init_ret value + 1
    switch_id = sys.argv[1]

    num_interfaces = wrapper.init(sys.argv[2:])
    interfaces = range(0, num_interfaces)

    MAC_table = {}
    
    switch_priority, port_vlan = get_ports_data(switch_id)
    bpdu = BPDU(switch_priority, port_vlan, interfaces)
    print("# Starting switch with id {}".format(switch_id), flush=True)
    print("[INFO] Switch MAC", ':'.join(f'{b:02x}' for b in get_switch_mac()))

    # Create and start a new thread that deals with sending BDPU
    t = threading.Thread(target=send_bdpu_every_sec, args=(bpdu,))
    t.start()

    # Printing interface names
    for i in interfaces:
        print(get_interface_name(i))

    while True:
        # Note that data is of type bytes([...]).
        # b1 = bytes([72, 101, 108, 108, 111])  # "Hello"
        # b2 = bytes([32, 87, 111, 114, 108, 100])  # " World"
        # b3 = b1[0:2] + b[3:4].
        interface, data, length = recv_from_any_link()

        dest_mac, src_mac, ethertype, vlan_id = parse_ethernet_header(data)

        # Print the MAC src and MAC dst in human readable format
        dest_mac = ':'.join(f'{b:02x}' for b in dest_mac)
        src_mac = ':'.join(f'{b:02x}' for b in src_mac)

        # Note. Adding a VLAN tag can be as easy as
        # tagged_frame = data[0:12] + create_vlan_tag(10) + data[12:]

        print(f'Destination MAC: {dest_mac}')
        print(f'Source MAC: {src_mac}')
        print(f'EtherType: {ethertype}')

        print("Received frame of size {} on interface {}".format(length, interface), flush=True)

        # TODO: Implement forwarding with learning
        MAC_table[src_mac] = interface
        if vlan_id != -1:
            vlan_in = vlan_id
        else:
            vlan_value = port_vlan.get(get_interface_name(interface), -1)
            if vlan_value == "T":
                vlan_in = -1  # Assign a default or meaningful value for trunk ports
            else:
                vlan_in = int(vlan_value)

        if bpdu.is_bpdu(interface, dest_mac, data, length):
            continue

        if (is_unicast(dest_mac) and
            dest_mac in MAC_table and
            MAC_table[dest_mac] in interfaces
    ):
                send_vlan_packet(MAC_table[dest_mac], data, length, port_vlan, vlan_id, vlan_in)
             
        else:
            for port in interfaces:
                if port != interface:
                    if bpdu.port_states[port] != "BLOCKING":
                        send_vlan_packet(port, data, length, port_vlan, vlan_id, vlan_in)

        # TODO: Implement VLAN support
        
        # TODO: Implement STP support

        # data is of type bytes.
        # send_to_link(i, length, data)

if __name__ == "__main__":
    main()
my stp algorithm is not correct implemented.Too many ICMPs arrived at the destination. Meaning we have a cycle.