#!/usr/bin/python -u
import sys
import socket
import time
import datetime
import select
import json

MSG_SIZE = 1500
DATA_SIZE = 1300
TIMEOUT = 30

# sequence is initialized at 0 but
# incremented based on length of data
# this number will represent the sequence
# number of the next acked packet we expect
SEQUENCE = 0

# number to represnt the last packet sents
# sequence number 
LASTSENTPACKETSEQNUM = 0

# constants for unacked buffer entries
# DATATABLE_SEQ | DATATABLE_DATA | DATATABLE_TIMESENT | DATATABLE_RESENT
# sequence number | data | time | resent or not
DATATABLE_SEQ = 0
DATATABLE_DATA = 1
DATATABLE_TIMESENT = 2
DATATABLE_RESENT = 3

# for the unAckedBuffer table, set the max size to deal with congestion
# this is bounded by the recievers advertized window 
# adv_wnd: receiver advertised window
# *** note for later. Don't know what this is right now
# will figure it out later
ADV_WND = 15

# cwnd:  congestion window
# CWND = min(CWND, ADV_WIND)
# but it is initialized to 1
CWND = 5

# ssthresh:  threshold size (used to update cwnd)
# but initially ssthresh = adv_wnd;
SSTHRESH = ADV_WND

# set default retransmittion timeout to 3 seconds
RTO = 2

# set round time trip alpha for recalculating
# the retransmission timeout
RTT_ALPHA = .8

# 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(RTO)

# Write message to stderr
def log(string):
  sys.stderr.write(datetime.datetime.now().strftime("%H:%M:%S.%f") + " " + string + "\n")

def send_next_packet():
  global timeLastPacketSent
  global LASTSENTPACKETSEQNUM
  global SEQUENCE

  # read in data from stdin, read in data size
  data = sys.stdin.read(DATA_SIZE)
  
  # if there is data, then create a message to send
  # if successfully sent, add to unacked packet buffer
  if (len(data) > 0):
    msg = json.dumps({"sequence": LASTSENTPACKETSEQNUM, "data": data, "ack": False, "eof": False})
    
    # initialize sequence num
    if (SEQUENCE == 0):
      SEQUENCE += len(data)
    
    # increment the sequence number so we can send
    # next packet after receiving right ack
    LASTSENTPACKETSEQNUM += len(data)
  
    if sock.sendto(msg, dest) < len(msg):
      log("[error] unable to fully send packet")
    else:
      # logs [send data] sequence_number (size of packet sent)
      log("[send data] " + str(LASTSENTPACKETSEQNUM) + " (" + str(len(data)) + ")")
      
      # update time last packet was sent
      timeLastPacketSent = time.time()
      
      # keep track of the packet we just sent in the unacked buffer table
      # sequence #, data, time sent, false because this is the first time
      # it is being sent
      addToUnackedBuffer(LASTSENTPACKETSEQNUM, data, timeLastPacketSent, False)
    return True
  else:
    return False
  
# send the last packet to the receiver
def sendFinalPacket():
  global LASTSENTPACKETSEQNUM
  global timeLastPacketSent
  
  sock.sendto(json.dumps({"eof": True, "data": "", "sequence": LASTSENTPACKETSEQNUM, "ack": False}), dest)
  
  # add 1 because ack num will be seq + 1
  LASTSENTPACKETSEQNUM = LASTSENTPACKETSEQNUM + 1
  
  # log it, update time of last sent packet and add to unacked buffer
  log("[send data] " + str(LASTSENTPACKETSEQNUM) + " (0)")
  timeLastPacketSent = time.time()
  addToUnackedBuffer(LASTSENTPACKETSEQNUM, "", timeLastPacketSent, False)
  
# keep track of all unacked packets  
def addToUnackedBuffer(seqNum, data, timeSent, resent):
  global unAckedBuffer
  
  # if already in table, update the time
  for entry in unAckedBuffer:
    # seq number could already be in the table
    # for 2 reasons
    # 1. we had to resend the packet
    # 2. we ran out of sequence numbers and it wrapped
    if entry[DATATABLE_SEQ] == seqNum:
      # note for future: may have to add more here
      entry[DATATABLE_DATA] = data
      entry[DATATABLE_TIMESENT] = timeSent
      entry[DATATABLE_RESENT] = resent
      return
  
  # otherwise add it
  unAckedBuffer.append([seqNum, data, timeSent, resent])
 
# update the sequence number when we get an
# ack in order
def updateSequenceNumber():
  global SEQUENCE
  
  #minSequenceNum = -1
  #for entry in unAckedBuffer:
  #  if minSequenceNum == -1 or entry[DATATABLE_SEQ] < minSequenceNum:
  #    minSequenceNum = entry[DATATABLE_SEQ]
      
  #SEQUENCE = minSequenceNum
  if len(unAckedBuffer) > 0:
    SEQUENCE = unAckedBuffer[0][DATATABLE_SEQ]      
 
# implement slow start or congestion avoidance, update CWND accordingly 
def updateCWND():
  global CWND
  global SSTHRESH
  
  if (CWND < SSTHRESH):
    # implement slow start
    CWND = CWND + 1
  else:
    #implement congestion avoidance
    CWND = CWND + 1/CWND

# update the congestion for when we hit the ACK timeout 
def updateCWNDonTimeout():
  global CWND
  global SSTHRESH
  
  #implement multiplicative decrease
  SSTHRESH = CWND/2
  CWND = 1;  

# once an ack is received for a packet
# remove it from the unacked buffer
def removeDataFromUnackedBuffer(ackNumber):
  for entry in unAckedBuffer:
    # check if this is a new ack
    if entry[DATATABLE_SEQ] == ackNumber:
      # if so remove the entry and update CWND
      unAckedBuffer.remove(entry)
      updateCWND()
      return
      
# RTO is a moving average so
# recalculate whenever we send a
# packet and receive an ack without
# having to retransmit
def calculateRTO(ackNumber):
  global RTO
  global sock
  global unAckedBuffer
  global RTT_ALPHA
  
  currentTime = time.time()
  
  for entry in unAckedBuffer:
    # find the entry for the ack we just received and use it to
    # calculate a new sample if it did not have to be retransmitted
    if entry[DATATABLE_SEQ] == ackNumber and not entry[DATATABLE_RESENT]:
      newSample = currentTime - entry[DATATABLE_TIMESENT]
      newRTT = (RTT_ALPHA * (RTO / 2)) + ((1 - RTT_ALPHA) * newSample)
      RTO = newRTT * 2
      sock.settimeout(RTO)
      log("New RTO " + str(RTO))
  
# check to see if any packets were
# dropped by checking ack timeouts    
def checkForAckTimeouts():
  global timeLastPacketSent
  
  currentTime = time.time()
  
  for entry in unAckedBuffer:
    # if packet dropped, resend
    # also update the congestion window
    if currentTime - entry[DATATABLE_TIMESENT] >= RTO:
      # update congestion window
      updateCWNDonTimeout()
      
      # sequence number is retrieved by getting current minus
      # the amount of data sent
      sequenceNum = entry[DATATABLE_SEQ] - len(entry[DATATABLE_DATA])
      
      # if packet was not the end of data, send it
      if len(entry[DATATABLE_DATA]) > 0:
        msg = json.dumps({"sequence": sequenceNum, "data": entry[DATATABLE_DATA], "ack": False, "eof": False})
      # if it was, set eof to true and send
      elif len(entry[DATATABLE_DATA]) == 0:
        msg = json.dumps({"sequence": sequenceNum - 1, "data": "", "ack": False, "eof": True})
       
      if sock.sendto(msg, dest) < len(msg):
        log("[error] unable to fully send packet")
      else:
        log("[send data] " + str(entry[DATATABLE_SEQ]) + " (" + str(len(entry[DATATABLE_DATA])) + ")")
        timeLastPacketSent = time.time()
        addToUnackedBuffer(entry[DATATABLE_SEQ], entry[DATATABLE_DATA], timeLastPacketSent, True)
    
      return
  
# send as many packets as CWND
# allows to go unacked      
def sendDataPackets():
  global unAckedBuffer
  global CWND
  global finalPacketSent
  
  diff = CWND - len(unAckedBuffer)
  
  while (diff > 0):
    if send_next_packet():
      diff = CWND - len(unAckedBuffer)
    elif not finalPacketSent:
      sendFinalPacket()
      finalPacketSent = 1
      return
    elif finalPacketSent:
      return

# keep track of unacked packets
# seq num, data, time sent
unAckedBuffer = []

# keep track of when last packet sent
# and received for timeout purposes
timeLastPacketSent = -1
timeLastPacketReceived = -1

# initially set finalPacketSent to false
finalPacketSent = 0

# Send first packet of data
send_next_packet()

# Now read in data, send packets
while True:
  log("ABOUT TO SLEEP")  
  try:
    result = sock.recvfrom(MSG_SIZE)

    if result:
      (data, addr) = result
      timeLastPacketReceived = time.time()
      try:
        decoded = json.loads(data)
        ackNumber = decoded['ack']
        
        #log("Ack received " + str(ackNumber))
	#log("Sequence num " + str(SEQUENCE))
        
        # If there is an ack, send more
        # packets and update the retransmission
        # timeout
        if ackNumber == SEQUENCE:
          calculateRTO(ackNumber)
          removeDataFromUnackedBuffer(ackNumber)
          log("[recv ack] " + str(ackNumber))
	     
          # Send as many data packets as the
          # congestion window allows
          sendDataPackets()
  
          # update sequence num        
          updateSequenceNumber()          
          
        # if the final packet has been sent and received
        # break out of while loop
        if finalPacketSent and len(unAckedBuffer) == 0:
          break
          
        # always check to see if packets
        # need to be resent
        checkForAckTimeouts()
        
      except (ValueError, KeyError, TypeError):
        log("[recv corrupt packet]")
  except socket.timeout:    
    # check to end the program
    # or if the message we are waiting
    # for an ack from timed out
    if timeLastPacketReceived != -1 and time.time() - timeLastPacketReceived >= TIMEOUT:    
      log("[error] timeout")
      sys.exit(-1)
    elif timeLastPacketSent != -1 and time.time() - timeLastPacketSent >= RTO:
      checkForAckTimeouts()

log("[completed]")
sys.exit(0)

# personal notes
 
 # sender side deals with congestion
 # window size determines send rate
 # adjust window to prevent congestion collapse
 
 # control the number of unACKed packets = adjust the size of the buffer 
 
 # wnd = min(cwnd, adv_wnd);
 # effective_wnd = wnd . (last_byte_sent . last_byte_acked);
 
 
 # Upon receipt of ACK: increase cwnd
 # Data was delivered, perhaps we can send faster
 # cwnd growth is proportional to RTT
 # On loss: decrease cwnd
 # Data is being lost, there must be congestion
 # Question: increase/decrease functions to use?
 
 # Maintains three variables:
 # cwnd:  congestion window
 # adv_wnd: receiver advertised window 
 # ssthresh:  threshold size (used to update cwnd)
 # For sending, use: cwnd = min(cwnd, adv_wnd)
 # Two phases of congestion control
 # Slow start (cwnd < ssthresh)
 # Probe for bottleneck bandwidth
 # Congestion avoidance (cwnd >= ssthresh)
 # AIMD
 
 # Goal: reach knee quickly
 # Upon starting (or restarting) a connection
 # cwnd =1
 # ssthresh = adv_wnd
 # Each time a segment is ACKed, cwnd++
 # Continues until.
 # ssthresh is reached
 # Or a packet is lost
 # Slow Start is not actually slow
 # cwnd increases exponentially
 
 # cwnd grows rapidly
 # Slows down when.
 # cwnd >= ssthresh
 # Or a packet drops
 
 # AIMD mode
 # ssthresh is lower-bound guess about location of the knee
 # If cwnd >= ssthresh then 
 #   each time a segment is ACKed cwnd += 1/cwnd
 # So cwnd is increased by one only if all segments have been acknowledged
 
 # Initially:
 #   cwnd = 1;
 #   ssthresh = adv_wnd;
 # New ack received:
 #   if (cwnd < ssthresh) 
 #         /* Slow Start*/
 #         cwnd = cwnd + 1;
 #   else
 #         /* Congestion Avoidance */
 #         cwnd = cwnd + 1/cwnd;
 # Timeout:
 #   /* Multiplicative decrease */
 #   ssthresh = cwnd/2;
 #   cwnd = 1;
 
 
 # Detect congestion
 # Packet dropping is most reliably signal
 # Delay-based methods are hard and risky
 # How do you detect packet drops? ACKs
 # Timeout after not receiving an ACK
 # Several duplicate ACKs in a row (ignore for now)
 # Rate adjustment algorithm
 # Modify cwnd
 # Probe for bandwidth
 # Responding to congestion


