#!/usr/bin/env python
from socket import socket, AF_INET, SOCK_DGRAM
from zlib import decompress, error, crc32
import datetime
import json
import sys
import random
#import select
#import time
UDP_IP = '127.0.0.1'
UDP_PORT = 0
MSG_SIZE = 1500
TIMEOUT = 30
SEQUENCE = random.randrange(4294967295) #random sequence number within the 32-bit range
AWINDOW = 10 #advertised window

msg_buff = {} #{ seq #, data }
latest_seq_in_order = 0
def log(string):
  sys.stderr.write(datetime.datetime.now().strftime("%H:%M:%S.%f") + " " + string + "\n")

def check_buff():
    global latest_seq_in_order
    global msg_buff
    # update the latest received packet in order by finding the first
    for seq in sorted(msg_buff)[1:]:
        if (seq-len(msg_buff[seq])) not in msg_buff:
            break
        else:
            length = len(msg_buff[seq])
            data = msg_buff[seq]
            log("[PRINTING FROM BUFF] " + str(seq) + " (" + str(len(data)) + ")")
            sys.stdout.write(data)
            # also manage the buffer size by deleting any that have already been written
            msg_buff.pop(latest_seq_in_order)
            # and the latest can only move by having the next sequence packet
            latest_seq_in_order = seq

if __name__=='__main__':
    # Set up the socket
    sock = socket(AF_INET, SOCK_DGRAM)
    sock.bind((UDP_IP,UDP_PORT))
    sock.settimeout(TIMEOUT)

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

    while True:
        result = sock.recvfrom(MSG_SIZE)
        # If nothing is ready, we hit the timeout
        if result:
            (data, addr) = result
        try:
            # the data may not always be compressed as there isn't always a benefit to it
            try:
                data = decompress(data)
            except error:
                log('receiving uncompressed data')
            decoded = json.loads(data)

            # If the CRC fails
            if 'crc' in decoded:
                recv_crc = int(decoded.pop('crc'))
                calc_crc = sum([crc32(str(val)) for key,val in decoded.iteritems()])
                if recv_crc != calc_crc:
                    log("crc mismatch: {} received and {} generated".format(recv_crc,calc_crc))
                    msg = json.dumps({"nack": decoded['sequence']})
                    log("ABOUT TO SEND " + msg)
                    if sock.sendto(msg, addr) < len(msg):
                        log("[error] unable to fully send packet")
                    continue

            if (decoded.get('syn')):
                # handle for third part of handshake
                if decoded.get('ack'):
                    msg_buff[decoded['sequence']+1] = ''
                    latest_seq_in_order = decoded['sequence']+1
                    log('[handshake] complete at SEQ#{}'.format(latest_seq_in_order))
                    continue

                # second message of handshake
                msg = json.dumps({"syn":decoded["sequence"],"ack": decoded['sequence']+1, "adv_window": AWINDOW})
            # If the EOF flag is set, exit
            elif (decoded.get('eof')):
                check_buff()
                if len(msg_buff)>1:
                    # ask for the next message if we didn't get the end
                    msg = json.dumps({"nack": latest_seq_in_order+len(msg_buff[latest_seq_in_order])})
                    sock.sendto(msg, addr)
                    continue

                msg = json.dumps({"fin": True})
                sock.sendto(msg,addr)
                log("[completed]")
                sys.exit(0)
            # If there is data, we accept it and print it out
            elif (decoded.get('data')):
                seq = decoded.get('sequence')
                if latest_seq_in_order == 0:
                    # we never successfully got the third part of the handshake before starting to get data
                    msg = json.dumps({"syn":False})
                #elif latest_seq_in_order != seq - len(decoded['data']):
                #    msg = json.dumps({"nack": latest_seq_in_order+len(msg_buff[latest_seq_in_order])})
                else:
                    # Send back an ack to the sender
                    msg = json.dumps({"ack": decoded['sequence'], "sequence": seq, "adv_window": AWINDOW})
                    # ignore duplicates
                    # TODO watch out for wraparound here.. also implement wraparound??
                    if not (seq in msg_buff or seq<= latest_seq_in_order):
                        msg_buff[decoded['sequence']]=decoded['data']
                        check_buff()
            log("ABOUT TO SEND " + msg)
            if sock.sendto(msg, addr) < len(msg):
                log("[error] unable to fully send packet")

        except (ValueError, KeyError, TypeError) as e:
            log("[recv corrupt packet]")
            log(str(e))
            msg = json.dumps({"nack": latest_seq_in_order+len(msg_buff[latest_seq_in_order])})
            if sock.sendto(msg, addr) < len(msg):
                raise e
    else:
        log("[error] timeout")
        sys.exit(-1)
