#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <poll.h>
#include <unistd.h>

#define MAX_CLIENTS 512
#define BACKLOG 10
#define DIE(info) ({perror((info)); exit(EXIT_FAILURE);})

static const char *socket_path = "/tmp/plumb-unix0";

struct plumb_client {
    int fd;
};

struct plumb_server {
    int socket_fd;
    int fd_num;
    struct plumb_client clients[MAX_CLIENTS];
    struct pollfd fds[MAX_CLIENTS + 1];
};


int plumb_server_client_add(struct plumb_server *s, int fd)
{
    printf("adding new connection, fd=%d\n", fd);

    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (s->clients[i].fd < 0) {
            s->clients[i].fd = fd;
            s->fds[i].fd = fd;
            s->fds[i].events = POLLIN;
            return i;
        }
    }
    return -1;
}

void plumb_server_client_remove(struct plumb_server *s, int i)
{
    if (i >= MAX_CLIENTS)
        return;

    printf("removing client i=%d\n", i);

    if (s->clients[i].fd >= 0) {
        close(s->clients[i].fd);
    }
    s->clients[i].fd = -1;
    s->fds[i].fd = -1;
    s->fds[i].events = 0;
    s->fds[i].revents = 0;
}


void plumb_server_free(struct plumb_server *s)
{
    if (!s) return;
    if (s->socket_fd) close(s->socket_fd);
    free(s);
}

struct plumb_server *plumb_server_create()
{
    struct plumb_server *s = malloc(sizeof(struct plumb_server));
    if (s == NULL) {
        DIE("malloc");
        return NULL;
    }

    s->socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (s->socket_fd == -1) {
        free(s);
        DIE("socket");
        return NULL;
    }

    struct sockaddr_un name;
    memset(&name, 0, sizeof(name));
    name.sun_family = AF_UNIX;
    strncpy(name.sun_path, socket_path, sizeof(name.sun_path) - 1);

    if (bind(s->socket_fd, (const struct sockaddr *)&name, sizeof(name)) == -1) {
        free(s);
        close(s->socket_fd);
        DIE("bind");
        return NULL;
    }
    if (listen(s->socket_fd, BACKLOG) == -1) {
        free(s);
        close(s->socket_fd);
        DIE("listen");
        return NULL;
    }

    for (int i = 0; i < MAX_CLIENTS; i++) {
        s->clients[i].fd = -1;
        s->fds[i].fd = -1;
        s->fds[i].events = 0;
        s->fds[i].revents = 0;
    }

    s->fds[MAX_CLIENTS].fd = s->socket_fd;
    s->fds[MAX_CLIENTS].events = POLLIN;

    s->fd_num = MAX_CLIENTS + 1;

    return s;
}

int plumb_server_turn(struct plumb_server *s)
{
    if (poll(s->fds, s->fd_num, -1) < 0) {
        DIE("poll");
        return -1;
    }

    for (int i = 0; i < s->fd_num; i++) {
        if (s->fds[i].revents & POLLIN) { // file descriptor ready for reading
            if (s->fds[i].fd == s->socket_fd) { // new connection
                int fd = accept(s->socket_fd, NULL, NULL);
                if (fd == -1) {
                    DIE("accept");
                    return -1;
                }

                if (plumb_server_client_add(s, fd) == -1) {
                    close(fd);
                    continue;
                }
            } else { // data from socket
                char data[256];
                memset(data, 0, sizeof(data));
                ssize_t bytes = recv(s->fds[i].fd, data, sizeof(data), 0);
                if (bytes == -1) {
                    DIE("recv");
                    return -1;
                } else if (bytes == 0) { // client disconnected
                    plumb_server_client_remove(s, i);
                } else {
                    send(s->fds[i].fd, data, sizeof(data), 0);
                }
            }
        }
    }

    return 0;
}

int main()
{
    struct plumb_server *srv = plumb_server_create();
    if (srv == NULL) {
        fprintf(stderr, "plumb_server_create failed\n");
        return EXIT_FAILURE;
    }

    while (1) {
        if (plumb_server_turn(srv) < 0) {
            fprintf(stderr, "plumb_server_turn failed\n");
            plumb_server_free(srv);
            return EXIT_FAILURE;
        }
    }

    plumb_server_free(srv);

    return EXIT_SUCCESS;
}