timeout = 1
serverAddressPort = ("213.47.107.152", 1337)

possible_value = input("Enter a value: ")
filename = possible_value if (filename := possible_value) else "data.bin"

import socket, struct, select, time, math

import hashlib

#import brotli #pip install brotlipy

max_frame_payload=508
header = "!L16s"
header_size = struct.calcsize(header)
max_payload = max_frame_payload - header_size

UDPClientSocket = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)


message = filename


UDPClientSocket.sendto(str.encode(message), serverAddressPort)

msgFromServer = UDPClientSocket.recvfrom(max_frame_payload)[0]

msg = struct.unpack(f'!LH{len(message)}s', msgFromServer)

last_packet_size = msg[1]

if msg[2] != str.encode(message):
    print('Server error: wrong file name')
    exit()

total_packets = msg[0]

packets = msg[0]

print(f'Packets to receive: {packets}')

total_data = b''

data_packets = []
for i in range(packets):
    data_packets.append(None)

sha512 = hashlib.sha512()

serverhash = None

time_started = time.time()

while packets > 0 or serverhash == None:

    ready = select.select([UDPClientSocket], [], [], timeout)

    if ready[0]:

        datareceived = UDPClientSocket.recvfrom(max_frame_payload)[0]
        if len(datareceived) == 64:
            serverhash = struct.unpack("64s", datareceived)[0]
            continue
        if len(datareceived) != max_frame_payload:
            print("\nerror bad packet: wrong size",len(datareceived))
            continue


        seq,checksum,data = struct.unpack(f"{header}{max_payload}s", datareceived)
        if data_packets[seq] != None:
            print("\nerror bad packet: received multiple times",seq)
            continue

        #check if packet is last packet and cut data to last packet size
        if seq == total_packets-1:
            #print("len before cut",len(data))
            data = data[:last_packet_size]
            # print("len after cut",len(data))

        #check if packet is corrupted
        if checksum.hex() != hashlib.sha3_512(data).digest()[:16].hex():
            print("\nerror bad packet: checksum mismatch",seq)
            print(checksum.hex(),hashlib.sha3_512(data).digest()[:16].hex())
            print(data.hex())
            continue

        data_packets[seq] = data

        packets -= 1
    else:
        #collect packets that were not received and send them in one message
        for i in range(len(data_packets)):
            if data_packets[i] == None:
                #print("requesting packet", i)
                UDPClientSocket.sendto(str.encode(f"{message}/{i}"), serverAddressPort)

        if serverhash == None:
            UDPClientSocket.sendto(str.encode(f"{message}:"), serverAddressPort)
    print(f'{total_packets - packets}/{total_packets} | {math.floor((total_packets - packets)/(time.time()-time_started if (time.time()-time_started)!=0 else 1)*max_payload)}Bps'+'\t'*10, end='\r')



print("finished transfer", end='\r')

for packet in data_packets:
    sha512.update(packet)

sha512 = sha512.hexdigest()

print('sha512: %s' % sha512)

print('server sha512: %s' % serverhash.hex())

if sha512 != serverhash.hex():
    print('hash mismatch')
    with open("received/"+filename+".corrupt", 'wb') as f:
        f.write(total_data)
    exit()

total_data = b''.join(data_packets)

with open("received/"+filename, 'wb') as f:
    f.write(total_data)

print(f'{filename} generated with size {len(total_data)} bytes')