/*
 *  Receive TCP packets using C (multiple clients, one thread, poll() loop)
 *
 *  Author: Tim Pierson, Dartmouth CS60, Fall 2026
 *          Port of Python version by Claude Opus 5
 *
 *  build: gcc -Wall -Wextra -o receive_tcp_polling receive_tcp_polling.c
 *  run:   ./receive_tcp_polling
 *
 *  send data from netcat in several terminals at once (replace IP address)
 *  nc -t 127.0.0.1 9090
 *  then type messages
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <poll.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define HOST "0.0.0.0"   /* "127.0.0.1" for localhost only */
#define PORT 9090
#define BUFSIZE 1024
#define MAX_CLIENTS 64

/* one of these per connected client. in the threaded version this state
   lived in handle_client's local variables; here there is no per-client
   stack to hold it, so we write it out by hand */
struct client {
    char who[INET_ADDRSTRLEN + 8];
    size_t bytes_seen;
};

/* fds[0] is always the listening socket; fds[1..nfds-1] are clients.
   clients[i] holds the state for fds[i], so the two arrays stay in step */
static struct pollfd fds[MAX_CLIENTS + 1];
static struct client clients[MAX_CLIENTS + 1];
static int nfds = 1;

static int send_all(int fd, const char *buf, size_t len)
{
    size_t sent = 0;
    while (sent < len) {
        ssize_t n = send(fd, buf + sent, len - sent, 0);
        if (n < 0) {
            if (errno == EINTR) continue;
            return -1;
        }
        sent += (size_t)n;
    }
    return 0;
}

/* drop slot i by moving the last entry into it. O(1), but it reorders the
   arrays, which is why the scan loop below must not assume a fixed order */
static void remove_client(int i)
{
    printf("Disconnected: %s\n", clients[i].who);
    close(fds[i].fd);
    fds[i]     = fds[nfds - 1];
    clients[i] = clients[nfds - 1];
    nfds--;
}

/* a client socket is readable: either data arrived or the client hung up */
static void on_readable(int i)
{
    char buf[BUFSIZE];
    char reply[BUFSIZE + 32];

    ssize_t n = recv(fds[i].fd, buf, sizeof buf, 0);
    if (n == 0) {           /* peer sent FIN: orderly shutdown, not an error */
        remove_client(i);
        return;
    }
    if (n < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) return;  /* nothing there after all */
        perror("recv");
        remove_client(i);
        return;
    }

    clients[i].bytes_seen += (size_t)n;
    printf("Received from %s: %.*s\n", clients[i].who, (int)n, buf);

    int len = snprintf(reply, sizeof reply, "Message received: %.*s", (int)n, buf);
    if (send_all(fds[i].fd, reply, (size_t)len) < 0) {
        perror("send");
        remove_client(i);
    }
}

/* the listening socket is readable: a client is waiting in the accept queue */
static void on_acceptable(void)
{
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof client_addr;

    /* accept returns a new socket to the client and the client's ip/port */
    int client_fd = accept(fds[0].fd, (struct sockaddr *)&client_addr, &addr_len);
    if (client_fd < 0) {
        if (errno != EAGAIN && errno != EWOULDBLOCK) perror("accept");
        return;
    }

    if (nfds > MAX_CLIENTS) {
        fprintf(stderr, "Too many clients, rejecting\n");
        close(client_fd);
        return;
    }

    fcntl(client_fd, F_SETFL, O_NONBLOCK);   /* never block the one thread we have */

    char ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &client_addr.sin_addr, ip, sizeof ip);

    fds[nfds].fd      = client_fd;
    fds[nfds].events  = POLLIN;
    fds[nfds].revents = 0;
    snprintf(clients[nfds].who, sizeof clients[nfds].who, "%s:%d",
             ip, ntohs(client_addr.sin_port));
    clients[nfds].bytes_seen = 0;
    nfds++;

    printf("Connected by %s:%d\n", ip, ntohs(client_addr.sin_port));
}

int main(void)
{
    struct sockaddr_in server_addr;

    /* create TCP socket listening on ip address and port */
    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) { perror("socket"); return 1; }

    int yes = 1;
    setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof yes);

    memset(&server_addr, 0, sizeof server_addr);
    server_addr.sin_family = AF_INET;
    server_addr.sin_port   = htons(PORT);
    inet_pton(AF_INET, HOST, &server_addr.sin_addr);

    if (bind(server_fd, (struct sockaddr *)&server_addr, sizeof server_addr) < 0) {
        perror("bind"); return 1;
    }
    if (listen(server_fd, 128) < 0) { perror("listen"); return 1; }

    fcntl(server_fd, F_SETFL, O_NONBLOCK);

    /* the listening socket is just another watched descriptor */
    fds[0].fd     = server_fd;
    fds[0].events = POLLIN;
    printf("Server listening on %s:%d...\n", HOST, PORT);

    /* the event loop. poll() blocks until at least one descriptor is ready,
       then we scan the array to find which. note the cost is O(watched),
       not O(ready) -- that is the limitation epoll and kqueue fix */
    for (;;) {
        int ready = poll(fds, (nfds_t)nfds, -1);   /* -1 means wait forever */
        if (ready < 0) {
            if (errno == EINTR) continue;
            perror("poll");
            break;
        }

        /* walk backwards so remove_client's swap-with-last cannot move an
           unexamined entry into a slot we have already passed */
        for (int i = nfds - 1; i >= 0; i--) {
            if (!(fds[i].revents & (POLLIN | POLLHUP | POLLERR))) continue;
            if (i == 0) on_acceptable();
            else        on_readable(i);
        }
    }

    for (int i = 0; i < nfds; i++) close(fds[i].fd);
    return 0;
}