from socket import *
import pickle
import sys
import threading

# TCP/IP socket
IP = 'localhost'
port = 40005

tcp_server = socket(AF_INET, SOCK_STREAM)
tcp_server.bind((IP, port))

tcp_server.listen(128)

print('Listen at port:', port)

socks = []
bufsize = 4096

client_params_list = []

def handle():
    while True:
        for idx, client_socket in enumerate(socks):
            try:
                msg = client_socket.recv(bufsize)
                if not msg:
                    continue
            except Exception:
                continue

            sep_idx = msg.find(b'\x80')
            if sep_idx < 0:
                data_len = int(msg.decode('utf-8'))
                from_client_msg = []
                msg_len = 0
            else:
                data_len = int(msg[:sep_idx].decode('utf-8'))
                from_client_msg = [msg[sep_idx:]]
                msg_len = len(msg[sep_idx:])
            
            while True:
                try:
                    packet = client_socket.recv(bufsize)
                except Exception:
                    continue

                from_client_msg.append(packet)
                msg_len += len(packet)

                if msg_len >= data_len:
                    break
            
            msg = b''.join(from_client_msg)
            print('Receive data, length:', len(msg))

            client_params = pickle.loads(msg)
            client_params_list[idx] = client_params

        if not None in client_params_list and len(client_params_list) > 0:
            client_avg_params = client_params_list[0]
            for name in client_avg_params:
                for i in range(1, len(client_params_list)):
                    client_avg_params[name] += client_params_list[i][name]
                client_avg_params[name] /= len(client_params_list)

            send_data = pickle.dumps(client_avg_params)
            send_info = str(len(send_data)).encode('utf-8')

            for idx, client_socket in enumerate(socks):
                client_socket.send(send_info)
                
                send_len = 0
                while send_len < len(send_data):
                    try:
                        send_len += client_socket.send(send_data[send_len:])
                    except:
                        pass

                print('Send data, length:', len(send_data))
                
                client_params_list[idx] = None
            

t = threading.Thread(target=handle)

if __name__ == '__main__':
    t.start()
    print()
    while True:
        client_socket, clientAddr = tcp_server.accept()
        print('---------------------------------------------')
        print('Client ' + str(len(socks)) + ', connected by', clientAddr)
        print('---------------------------------------------')
        client_socket.setblocking(0)
        socks.append(client_socket)
        client_params_list.append(None)