"""
    Simulation model components for an input-queueing packet switch.
    To run, they need Dr. Bernstein's SimComponents.py 
    See: https://www.grotto-networking.com/DiscreteEventPython.html

    Leon Seng and Ahmet Sekercioglu, last modified on 20 May 2020
"""
import simpy
import random


class IQSwitchInputPort(object):
    """ Models a switch input port that demultiplexes incoming packets to random
        output ports at a given rate and buffer size limit in bytes.

        Contains a list of output ports of the same length as the
        probability list in the constructor. Use these to connect to
        other network elements.

        Packets are forwarded to randomly selected output ports when
        the ports are free. An output port is polled at a specified
        rate to determine if it is free.

        Parameters
        ----------
        env : simpy.Environment
            the simulation environment
        rate : float
            the bit rate of the port
        qlimit : integer (or None)
            a buffer size limit in bytes or packets for the queue 
            (including items in service).
        limit_bytes : If true, the queue limit will be based on 
            bytes if false the
            queue limit will be based on packets.
        probs : List
            list of probabilities for the corresponding output ports
    """
    def __init__(self, env, id, rate, probs, qlimit=None, limit_bytes=True, debug=False):
        self.store = simpy.Store(env)
        self.id = id
        self.rate = rate
        self.env = env
        self.packets_rec = 0
        self.packets_drop = 0
        self.qlimit = qlimit
        self.limit_bytes = limit_bytes
        self.byte_size = 0  # Current size of the queue in bytes
        self.debug = debug
        self.busy = 0  # Used to track if a packet is currently being sent

        # define outputs
        self.probs = probs
        self.ranges = [sum(probs[0:n+1]) for n in range(len(probs))]  # Partial sums of probs
        if self.ranges[-1] - 1.0 > 1.0e-6:
            raise Exception("Probabilities must sum to 1.0")
        self.n_ports = len(self.probs)
        self.outs = [None for i in range(self.n_ports)]  # Create and initialize output ports
        self.packets_rec = 0

        self.action = env.process(self.run())  # starts the run() method as a SimPy process

    def run(self):
        while True:
            msg = (yield self.store.get())

            # select output
            rand = random.random()
            output = None
            for i in range(self.n_ports):
                if rand < self.ranges[i]:
                    output = self.outs[i]
                    break

            # only forward msg if output is defined, else "drop"
            if not output:
                continue

            while not output.link_port(self.id):
                yield self.env.timeout(1/self.rate)

            self.busy = 1
            self.byte_size -= msg.size
            yield self.env.timeout(msg.size*8.0/self.rate)
            output.put(msg)
            output.unlink_port(self.id)
            self.busy = 0
            if self.debug:
                print(msg)

    def put(self, pkt):
        self.packets_rec += 1
        tmp_byte_count = self.byte_size + pkt.size

        if self.qlimit is None:
            self.byte_size = tmp_byte_count
            return self.store.put(pkt)
        if self.limit_bytes and tmp_byte_count >= self.qlimit:
            self.packets_drop += 1
            return
        elif not self.limit_bytes and len(self.store.items) >= self.qlimit-1:
            self.packets_drop += 1
        else:
            self.byte_size = tmp_byte_count
            return self.store.put(pkt)


class IQSwitchOutputPort(object):
    """ Models an NxN input queued switch output port with a given rate
        and buffer size limit in bytes.

        An input port needs to reserve this output port object before it is
        able to perform a put

        Set the "out" member variable to the entity to receive the packet.

        Parameters
        ----------
        env : simpy.Environment
            the simulation environment
        rate : float
            the bit rate of the port

    """
    def __init__(self, env, rate, debug=False):
        self.store = simpy.Store(env)
        self.rate = rate
        self.env = env
        self.out = None
        self.packets_rec = 0
        self.packets_drop = 0
        self.byte_size = 0  # Current size of the queue in bytes
        self.debug = debug
        self.busy = 0  # Used to track if a packet is currently being sent
        self.authorized_packet_source = None  # keep track of which input port is sending packet to this output port
        self.reserved = False
        self.action = env.process(self.run())  # starts the run() method as a SimPy process

    def run(self):
        while True:
            msg = (yield self.store.get())
            self.busy = 1
            self.byte_size -= msg.size
            yield self.env.timeout(msg.size*8.0/self.rate)
            self.out.put(msg)
            self.busy = 0
            if self.debug:
                print(msg)

    def link_port(self, input_id):
        """
        Reserve this output port
        """
        if (
            self.reserved or
            self.busy or
            len(self.store.items) > 0
        ):
            # print(f"Failed to link {input_id}: {self.reserved} {self.busy} {len(self.store.items)}")
            return False
        else:
            self.authorized_packet_source = input_id
            self.reserved = True
            # if self.debug:
            #     print(f"Linked {input_id}")
            return True

    def unlink_port(self, input_id):
        """
        Remove link between an input port to an output port
        """
        if self.authorized_packet_source == input_id:
            self.reserved = False
            self.authorized_packet_source = None
        else:
            raise Exception(f"Requester {input_id} is not currently the registered packet source {self.authorized_packet_source}")

    def put(self, pkt):
        if not self.reserved:
            raise Exception(f"Output port unreserved before putting packet.")

        self.packets_rec += 1
        tmp_byte_count = self.byte_size + pkt.size

        self.byte_size = tmp_byte_count
        return self.store.put(pkt)