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

def log(string):
  sys.stderr.write(datetime.datetime.now().strftime("%H:%M:%S.%f") + " " + string + "\n")
  
## set the ackNumber to the sequence num
def initializeAckNumber(packet_sequence):
  global ackNumber
  ackNumber = packet_sequence
  
# Sort the sequence nums in dictionary
# and print out received data
def printReceivedData():
  global receivedData
  
  for dataEntry in sorted(receivedData):
    sys.stdout.write(receivedData[dataEntry])
    
# send ack back to sender after receiving
# the data given
def sendAck(packetAckNumber):  
  msg = json.dumps({"ack": packetAckNumber})
  log("ABOUT TO SEND " + msg)
  
  # if sock.sendto() is less than full size of message
  # it means the full message didn't send
  #TODO: better checking here?
  if sock.sendto(msg, addr) < len(msg):
    log("[error] unable to fully send packet")
    
# increment ack num by data length
# or by 1 if it is the last packet
def incrementAck(packet_data):
  global ackNumber
  
  if len(packet_data) > 0:
    ackNumber += len(packet_data)
  else:
    ackNumber = packet_sequence + 1

# after getting final packet, send
# ack and exit
def closeReceiver():
  global ackNumber
  
  printReceivedData()
  incrementAck("")
  sendAck(ackNumber)
  log("[completed]")
  sys.exit(0)

# look to see if we have already received
# the next in order packet from the sender
# if we have, send an ack for it
def checkOutOfOrderPackets():
  global ackNumber

  check = True

  while check:
     if ackNumber in outOfOrderPackets:
       log("Sending ack from packet out of order")
       outOfOrderPackets.remove(ackNumber)
       incrementAck(receivedData[ackNumber])
       sendAck(ackNumber)
     else:
       check = False

# set global variables
MSG_SIZE = 1500
TIMEOUT = 30

# Bind to localhost and an ephemeral port
UDP_IP = "127.0.0.1"
UDP_PORT = 0

# Set up the socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) 
sock.bind((UDP_IP, UDP_PORT))

# timeout set to 30 seconds
sock.settimeout(TIMEOUT)

# Get port we bound to
UDP_PORT = sock.getsockname()[1]
log("[bound] " + str(UDP_PORT))

# what should the buffer size be set to?
# receivedData{} should have a maximum size if we 
# want to implement a sliding window later. Congestion 
# is delt with on the sender side 
# Store the sequence numbers and 
# corresponding data successfully received
receivedData = {}

# Keep track of the last ack number
# to see if next segment received is
# in order or not, start at -1
ackNumber = -1

# keep track of packets we receive but not
# in order so unacked
outOfOrderPackets = []

# Now listen for packets
while True:
  result = sock.recvfrom(MSG_SIZE)

  # If nothing is ready, we hit the timeout
  if result:
    (data, addr) = result
    
    try:
      # Get info out of packet
      decoded = json.loads(data)
      packet_eof = decoded['eof']
      packet_sequence = decoded['sequence']
      packet_data = decoded['data']
      
      # If it is the first packet we get, set
      # ack to sequence number
      if (ackNumber == -1):
        initializeAckNumber(packet_sequence)
      
      # if end of packets, print out data received
      # and close receiver
      #log(str(ackNumber) + " " + str(packet_sequence))
      if (packet_eof and ackNumber == packet_sequence):
        closeReceiver()
      
      # If there is data, look at it and 
      # handle accordingly
      if (packet_data):   
        
          # if we get a duplicated message this could be because it was dropped
          # check to see if ack needs to be incremented and send ack with
          # the correct ack number for this message
          if (packet_sequence in receivedData):
            log("[recv data] " + str(packet_sequence) + " (" + str(len(packet_data)) + ") IGNORED")
            sendAck(packet_sequence + len(packet_data))

	    if (ackNumber == packet_sequence):
	      incrementAck(packet_data)  
           
          else:
            # packet sent in order, send back ack with incremented ack number
            if (ackNumber == packet_sequence):
              incrementAck(packet_data)
              log("[recv data] " + str(packet_sequence) + " (" + str(len(packet_data)) + ") ACCEPTED (in-order)")
              sendAck(ackNumber)

	      checkOutOfOrderPackets()
            # packet sent out of order, send back ack with unchanged ack number
            elif (ackNumber != packet_sequence):
              log("[recv data] " + str(packet_sequence) + " (" + str(len(packet_data)) + ") ACCEPTED (out-of-order)")
	      outOfOrderPackets.append(packet_sequence)
              #sendAck(ackNumber)
        
            # store the data
            receivedData[packet_sequence] = packet_data


    except (ValueError, KeyError, TypeError) as e:
      log("[recv corrupt packet]")
      raise e
  else:
    log("[error] timeout")
    sys.exit(-1)

# personal notes
#
# on recieving end we can:
# set the advertised window
# ignore duplicate packets
# Out-of-order packets are reordered or dropped
# Missing sequence numbers indicate lost packets

