#!/usr/bin/env -S python3 -u

import argparse, socket, time, json, select, struct, sys, math, time, copy

DATA_SIZE = 1375

class Sender:
    def __init__(self, host, port):
        self.host = host
        self.remote_port = int(port)
        self.log("Sender starting up using port %s" % self.remote_port)
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(('0.0.0.0', 0))
        self.waiting = False

        self.seen_seq_no = set()
        self.seq_no = 0

        self.done_sending = False

        self.awaiting_ack = []
        self.retransmit = []

    def log(self, message):
        sys.stderr.write(message + "\n")
        sys.stderr.flush()

    def send(self, message):
        self.socket.sendto(json.dumps(message).encode('utf-8'), (self.host, self.remote_port))

    def run(self):
        while True:
            sockets = [self.socket, sys.stdin] if not self.waiting else [self.socket]

            socks = select.select(sockets, [], [], 0.1)[0]
            for conn in socks:
                if conn == self.socket:
                    # 3 packet acks can be received simultaneously
                    for i in range(3):
                        k, addr = conn.recvfrom(65535)
                        msg = k.decode('utf-8')

                        self.log("Received message '%s'" % msg)
                        ack_seqno = json.loads(msg)['seqno']
                        self.seen_seq_no.add(ack_seqno)

                        self.detect_drop(ack_seqno)
                        #self.log("post drop retransmit list "+str(self.retransmit))
                        #self.log("post drop awaiting_ack list "+str(self.awaiting_ack))

                        if self.done_sending:
                            if ack_seqno == self.seq_no:
                                self.log("final ack received")
                                sys.exit(0)
                        self.waiting = False
                    
                elif conn == sys.stdin:
                    # up to 3 packets can be retransmitted simultaneously
                    for p in range(min(3,len(self.retransmit))):
                        msg = { "type": "msg", "seqno": self.retransmit[p][0], "data": self.retransmit[p][1] }
                        self.log("Retransmitting message '%s'" % msg)
                        self.send(msg)
                    self.retransmit = self.retransmit[min(3,len(self.retransmit)):]

                    # 3 - # retransmissions can be sent
                    for i in range(3-min(3,len(self.retransmit))):
                        data = sys.stdin.read(DATA_SIZE)
                        if len(data) == 0:
                            self.done_sending = True
                            self.log("All done!")
                            if self.seq_no == max(self.seen_seq_no):
                                self.log("final ack received")
                                sys.exit(0)
                            break

                        msg = { "type": "msg", "seqno": self.seq_no, "data": data }
                        self.seen_seq_no.add(self.seq_no)
                        self.seq_no += 1
                        self.log("Sending message '%s'" % msg)
                        self.send(msg)

                        # save the time and details corresponding to each sent packet
                        self.awaiting_ack.append((time.time(),self.seq_no-1,data))
                    self.waiting = True

        return

    # detect if a packet has been dropped from the network
    # add packet to retransmission queue if the ack was not received in time
    def detect_drop(self,ack_no):
        temp_list = []
        # starttime, seqno, data
        for p in self.awaiting_ack:
            endtime = time.time()
            # ack not received in time - packet must be retransmitted
            '''self.log('packet'+str(p[1])+' sent at '+str(p[0]))
            self.log('ack'+str(ack_no)+'received at'+str(endtime))
            self.log('difference:'+str(endtime - p[0]))'''
            if endtime - p[0] >= 1.0:
                self.log('ADDED TO RETRANSMIT'+str(p[1]))
                self.retransmit.append((p[1],p[2]))
            elif p[1] == ack_no - 1 and endtime - p[0] <= 1.0:
                #self.log('REMOVING '+str(p[1])+' FROM AWAITING')
                continue
            # still awaiting ack
            elif endtime - p[0] <= 1.0:
                temp_list.append(p)
        self.awaiting_ack = temp_list

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='send data')
    parser.add_argument('host', type=str, help="Remote host to connect to")
    parser.add_argument('port', type=int, help="UDP port number to connect to")
    args = parser.parse_args()
    sender = Sender(args.host, args.port)
    sender.run()