#!/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

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, lenght, data, 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():
    while True:
        # 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)

    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)
    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_in = int(port_vlan[get_interface_name(interface)])

        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], length, data, port_vlan, vlan_id, vlan_in)
             
        else:
            for port in interfaces:
                if port != interface:
                    send_vlan_packet(port, length, data, 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()

Traceback (most recent call last):
  File "/home/traistaru-andreea/Desktop/RL/Tema_1/tema1-public/switch.py", line 153, in <module>
    main()
  File "/home/traistaru-andreea/Desktop/RL/Tema_1/tema1-public/switch.py", line 143, in main
    send_vlan_packet(port, length, data, port_vlan, vlan_id, vlan_in)
  File "/home/traistaru-andreea/Desktop/RL/Tema_1/tema1-public/switch.py", line 23, in send_vlan_packet
    send_to_link(interface, length, data)
                            ^^^^^^
UnboundLocalError: cannot access local variable 'length' where it is not associated with a value