import socket
import selectors

''' 
    Receive TCP packets using Python (multiple clients, one thread, polling loop)

	Author: Tim Pierson, Dartmouth CS60, Fall 2026
            Assisted by Microsoft Copilot
	
	run: python3 receive_tcp_polling.py
	
	send data from netcat in several terminals at once (replace IP address)
	nc -t 127.0.0.1 9090
	then type messages
'''

#define host ip and port
HOST = '0.0.0.0' #'127.0.0.1' #localhost
PORT = 9090

#the selector is the kernel's "tell me which sockets are ready" service.
#DefaultSelector picks epoll on linux, kqueue on mac, select on windows
sel = selectors.DefaultSelector()

#one of these per connected client. in the threaded version this state lived
#in the handler function's local variables; here there is no per-client stack
#to hold it, so we write it out by hand
class Client:
    def __init__(self, client_socket, addr):
        self.socket = client_socket
        self.addr = addr
        self.bytes_seen = 0

#a client socket is readable: either data arrived or the client hung up
def on_readable(client):
    data = client.socket.recv(1024)
    if not data: #empty bytes means the client sent FIN, not "no data yet"
        print(f"Disconnected: {client.addr}")
        sel.unregister(client.socket) #stop watching before closing
        client.socket.close()
        return
    client.bytes_seen += len(data)
    print(f"Received from {client.addr}: {data.decode()}") #decode converts bytes to string
    client.socket.sendall(b"Message received: " + data) #sendall sends all chunks (send only sends one chunk!)

#the listening socket is readable: a new client is waiting in the accept queue
def on_acceptable(server_socket):
    client_socket, addr = server_socket.accept() #accept returns a new socket to the client and the clients ip/port
    print(f"Connected by {addr}")
    client_socket.setblocking(False) #never block the one thread we have
    client = Client(client_socket, addr)
    sel.register(client_socket, selectors.EVENT_READ, lambda: on_readable(client))

#create TCP socket listening on ip address and port
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) #lets us restart right away instead of waiting out TIME_WAIT
server_socket.bind((HOST, PORT))

#listen for connections
server_socket.listen()
server_socket.setblocking(False)
sel.register(server_socket, selectors.EVENT_READ, lambda: on_acceptable(server_socket))
print(f"Server listening on {HOST}:{PORT}...")

#the event loop. select() blocks until at least one socket is ready, then hands
#back only those, and we run whatever callback we registered for each
try:
    while True:
        for key, mask in sel.select():
            key.data() #key.data is the lambda we registered above
except KeyboardInterrupt:
    print("\nShutting down.")
finally:
    sel.close()
    server_socket.close()