#!/usr/bin/python3 -u
# Usage: ./3700send <recv_host>:<recv_port>
import sys
import struct
import socket
import time
import datetime
import select

MSG_SIZE = 1500
DATA_SIZE = 1000
TIMEOUT = 5
SEQUENCE = 0

# global storage, key = expected ACK if corresponding datagram is received
DATA = {}
SAVED_ACKS = []
# timestamps, only updated when a NEW transmission is sent
send_time = datetime.datetime.now()
rec_time = datetime.datetime.now()
last_sent = datetime.datetime.now()
RTT = 0  # round trip time
RTO = 0  # retransmission time out
CWND = 10  # congestion window size
ADV_WND = 65535  # maximum size of TCP buffer
SSTHRESH = ADV_WND  # slow start threshold

# Bind to localhost and an ephemeral port
IP_PORT = sys.argv[1]
UDP_IP = IP_PORT[0:IP_PORT.find(":")]
UDP_PORT = int(IP_PORT[IP_PORT.find(":") + 1:])
dest = (UDP_IP, UDP_PORT)

# Set up the socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.settimeout(TIMEOUT)

# communication channels
inputs = [sock]
outputs = []


# calculate new sample of round trip time to adjust retransmission time out
def calculateRTT():
    global RTT
    new_sample = (rec_time - send_time).total_seconds()
    alpha = 0.875
    RTT = (alpha * RTT) + (1 - alpha) * new_sample


# update retransmission time out by first updated RTT then getting the
# RTO from that new RTT sample
# SHOULD ONLY BE CALLED ON PACKETS NOT BEING RETRANSMITTED
def updateRTO():
    global RTO, RTT
    calculateRTT()
    RTO = max(RTT * 2, 0.2)


# update congestion window based on difference of current ack and most recent ack
def updateCWND(curr_ack):
    global CWND
    CWND = curr_ack
    log("CWND: " + str(CWND))



# updates the sequence number so the sequence # equals
# the last received ack / the sequence # of the next datagram to send
def updateSequence(ack_num):
    global SEQUENCE
    SEQUENCE = ack_num


# add to data array when sending a new data package
def addToData(data):
    global SEQUENCE, DATA_SIZE
    # the ack expected if this data is received
    expectedAck = SEQUENCE + DATA_SIZE
    DATA[expectedAck] = data


# delete from data array once we received the ack for that data package
def deleteFromData(ack):
    global DATA
    if ack in DATA:
        del DATA[ack]


# retransmit a data package because sender timed out for an ACK response
def retransmit():
    global DATA, SEQUENCE, DATA_SIZE
    # get the data of the ACK that we were expecting (but didn't get)
    msg = DATA[SEQUENCE + DATA_SIZE]
    sock.sendto(msg, dest)


def filterSavedAcks():
    global SAVED_ACKS, SEQUENCE
    while SEQUENCE in SAVED_ACKS:
        SAVED_ACKS.remove(SEQUENCE)
        updateSequence(SEQUENCE + DATA_SIZE)


# print statement
def log(string):
    sys.stderr.write("[" + datetime.datetime.now().strftime("%H:%M:%S.%f") + "]" + " SENDER: " + string + "\n")


# send multiple packets at once
def send_multiple_packets():
  global CWND
  next_msg_ack = SEQUENCE + DATA_SIZE
  while next_msg_ack < CWND:
      log("Send next packet!")
      send_next_packet()


# sends the next packet to recv
def send_next_packet():
    global SEQUENCE, send_time, DATA_SIZE

    data = sys.stdin.buffer.read(DATA_SIZE)
    # there is more data to be sent to recv
    if len(data) > 0:
        # make header
        header_format = (SEQUENCE, False)
        header_struct = struct.Struct('I ?')
        header = header_struct.pack(*header_format)

        # make datagram
        msg = bytearray()
        msg.extend(header)
        msg.extend(data)

        # add msg to data dict
        addToData(msg)

        # send the datagram
        sock.sendto(msg, dest)

        # update the send_time to correspond to now
        # since we just sent a datagram
        send_time = datetime.datetime.now()
        log("NEXT PACKET HAS BEEN SENT")

        # this is not the last packet
        return True
    # there is no more data to be sent to recv, send final datagram indicating EOF
    else:
        # make header
        header_format = (SEQUENCE, True)
        header_struct = struct.Struct('I ?')
        header = header_struct.pack(*header_format)

        # make datagram
        msg = bytearray()
        msg.extend(header)

        # send the datagram
        sock.sendto(msg, dest)
        log("LAST PACKET HAS BEEN SENT")

        # this is the last packet
        return False


# send the first packet
send_next_packet()

# Now continuously read in data and send packets
while inputs:
    log("READING DATA...")
    curr_time = datetime.datetime.now()

    # how much time as elapsed since the last datagram was sent?
    secs_elapsed = (curr_time - last_sent).total_seconds()

    # checks if ack data can be received
    reads, writes, errors = select.select(inputs, outputs, inputs, RTO)

    # receive message from client socket
    try:
        result = sock.recvfrom(MSG_SIZE)
    except:
        result = None

    log('SECS ELAPSED: ' + str(secs_elapsed))
    # proceed if data is received from client
    if result:
        (received, addr) = result
        ack_num = int.from_bytes(received, byteorder='big')
        log("ACK_NUM: " + str(ack_num))
        log("SEQUENCE: " + str(SEQUENCE))
        # if ACK corresponds to the last sent datagram
        if SEQUENCE + DATA_SIZE == ack_num:
            log("THE STUFF IS NEXT IN ORDER")
            rec_time = datetime.datetime.now()

            # when new ack received, commence slow start if congestion window is less than slow start threshold, else do congestion avoidance
            if CWND < SSTHRESH:
                new_cwnd = CWND + 1
            else:
                new_cwnd = round((CWND + 1) / CWND)

            updateCWND(new_cwnd)
            updateRTO()
            deleteFromData(ack_num)
            updateSequence(ack_num)
            filterSavedAcks()

            if not send_next_packet():
                # sender is done sending
                log("[completed]")
                break
        elif SEQUENCE + DATA_SIZE < ack_num:
            log("ACKS ARE BEING RECEIVED OUT OF ORDER")
            # if (SEQUENCE < ack_num) the packets are being sent out of order
            SAVED_ACKS.append(ack_num)
            # update congestion window based on most recent ack
            new_cwnd = CWND + ack_num - (SEQUENCE + DATA_SIZE)
            log("NEW CWND: " + str(new_cwnd))
            updateCWND(new_cwnd)
            break

    # if we have NOT gotten a response from the client, AND it has been longer than RTO = timeout
    else:
        log('TIMEOUT OCCURRED')
        # multiplicative decrease
        SSTHRESH = round(CWND / 2)
        # reset congestion window to 1
        updateCWND(1)
        # retransmit the last sent datagram
        retransmit()
        log("RETRANSMITTED DATAGRAM")
        # change last_sent to be now since we just re-sent a datagram right now
        last_sent = datetime.datetime.now()
        send_multiple_packets()

sys.exit(0)