import socket
import threading
import struct
import traceback

def send_one_message(sock, data):
    length = len(data)
    sock.sendall(struct.pack('!I', length))
    sock.sendall(data)

def recv_one_message(sock):
    lengthbuf = recvall(sock, 4)
    length, = struct.unpack('!I', lengthbuf)
    return recvall(sock, length)

def recvall(sock, count):
    fragments = []
    while count: 
        chunk = sock.recv(count)
        if not chunk: 
            return None
        fragments.append(chunk)
        count -=len(chunk)
    arr = b''.join(fragments)
    return arr

def handle_client(client_socket, address):
    print(f'{socket.gethostname()} New client connected: {address}')
    try:
        while True:
            data = recv_one_message(client_socket)
            if not data:
                break
            message = data.decode().strip().lower() # convert message to lowercase
            print(f'{socket.gethostname()} Received message from {address}: {message}')
            if message == 'hello':
                response = 'Hello'
                send_one_message(client_socket,response.encode())
            elif message == 'file':
                # receive filename
                filename = recv_one_message(client_socket).decode().strip()
                print(f'{socket.gethostname()} Received filename from {address}: {filename}')

                # receive filesize
                filesize = int(recv_one_message(client_socket).decode().strip())
                print(f'{socket.gethostname()} Received filesize from {address}: {filesize}')

                # receive file data
                with open(filename, 'wb') as f:
                    remaining_bytes = filesize
                    while remaining_bytes > 0:
                        data = recv_one_message(client_socket)
                        if not data:
                            break
                        f.write(data)
                        remaining_bytes -= len(data)
                print(f'{socket.gethostname()} Received the entire file from {address}: {filesize}')
                response = 'File received'
                send_one_message(client_socket,response.encode())
            elif message == 'bye':
                break
    except Exception as e:
        print(f'{socket.gethostname()} Error handling client {address}: {e}')
        traceback.print_exc()
    finally:
        client_socket.close()
        print(f'{socket.gethostname()} Client disconnected: {address}')

def start_server(port):
    try:
        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_socket.bind(('0.0.0.0', port))
        server_socket.listen()
        print(f'{socket.gethostname()} Server listening on port {port}')

        while True:
            client_socket, address = server_socket.accept()
            print(f'{socket.gethostname()} Accepted new connection from {address}')
            client_thread = threading.Thread(target=handle_client, args=(client_socket, address))
            client_thread.start()
    except Exception as e:
        print(f'{socket.gethostname()} Error starting server: {e}')
        traceback.print_exc()
    finally:
        server_socket.close()

if __name__ == '__main__':
    port = 65122 # pick any free port you wish that is not used by other students
    start_server(port)