/*
 * 04_epoll_oneshot.c —— EPOLLONESHOT + 多线程 worker 模板
 *
 * 编译: gcc 04_epoll_oneshot.c -o oneshot_srv -pthread
 * 运行: ./oneshot_srv 9003
 * 测试: 同时多个 nc localhost 9003，会被多个 worker 线程公平消费
 *
 * 思路：主线程 epoll_wait → 把就绪的 fd 派给一个 worker 线程
 *      EPOLLONESHOT 保证同一个 fd 一次只被一个 worker 处理
 *      worker 处理完必须重新 EPOLL_CTL_MOD 把它"挂回去"
 */

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

#define MAX_EVENTS 64
#define WORKERS    4
#define BUF_SIZE   1024

static int epfd;

/* 简单的有锁队列 */
typedef struct {
    int items[1024];
    int head, tail, size;
    pthread_mutex_t mu;
    pthread_cond_t  cv;
} job_queue_t;

static job_queue_t Q;

static void q_init() {
    Q.head = Q.tail = Q.size = 0;
    pthread_mutex_init(&Q.mu, NULL);
    pthread_cond_init(&Q.cv, NULL);
}
static void q_push(int fd) {
    pthread_mutex_lock(&Q.mu);
    if (Q.size < 1024) {
        Q.items[Q.tail] = fd;
        Q.tail = (Q.tail + 1) % 1024;
        Q.size++;
        pthread_cond_signal(&Q.cv);
    }
    pthread_mutex_unlock(&Q.mu);
}
static int q_pop() {
    pthread_mutex_lock(&Q.mu);
    while (Q.size == 0) pthread_cond_wait(&Q.cv, &Q.mu);
    int fd = Q.items[Q.head];
    Q.head = (Q.head + 1) % 1024;
    Q.size--;
    pthread_mutex_unlock(&Q.mu);
    return fd;
}

static void set_nonblock(int fd) {
    int fl = fcntl(fd, F_GETFL, 0);
    fcntl(fd, F_SETFL, fl | O_NONBLOCK);
}

static void rearm(int fd) {
    /* 重新挂回 epoll，否则这个 fd 永远不会再触发 */
    struct epoll_event ev = {
        .events = EPOLLIN | EPOLLET | EPOLLONESHOT | EPOLLRDHUP,
        .data.fd = fd,
    };
    if (epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev) < 0) {
        if (errno != ENOENT) perror("EPOLL_CTL_MOD");
    }
}

static void *worker(void *arg) {
    long wid = (long)arg;
    while (1) {
        int fd = q_pop();
        printf("[worker-%ld] 接到任务 fd=%d\n", wid, fd);

        int closed = 0;
        while (1) {
            char buf[BUF_SIZE];
            ssize_t r = read(fd, buf, sizeof(buf));
            if (r > 0) {
                write(fd, buf, r);
                continue;
            }
            if (r == 0) { close(fd); closed = 1; break; }
            if (errno == EAGAIN || errno == EWOULDBLOCK) break;
            close(fd); closed = 1; break;
        }
        if (!closed) rearm(fd);    /* 处理完后必须重新挂回 */
    }
    return NULL;
}

static int make_listen_socket(int port) {
    int fd = socket(AF_INET, SOCK_STREAM, 0);
    int yes = 1; setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
    set_nonblock(fd);
    struct sockaddr_in addr = {0};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    if (bind(fd, (struct sockaddr*)&addr, sizeof(addr)) < 0 || listen(fd, 64) < 0) {
        perror("bind/listen"); exit(1);
    }
    printf("🚀 ONESHOT + 多线程 worker 服务器启动，端口 %d (workers=%d)\n", port, WORKERS);
    return fd;
}

int main(int argc, char *argv[]) {
    int port = (argc > 1) ? atoi(argv[1]) : 9003;
    int lfd  = make_listen_socket(port);
    epfd = epoll_create1(EPOLL_CLOEXEC);

    /* 监听 socket: 不加 ONESHOT，因为只有主线程操作它 */
    struct epoll_event ev = {.events = EPOLLIN | EPOLLET, .data.fd = lfd};
    epoll_ctl(epfd, EPOLL_CTL_ADD, lfd, &ev);

    q_init();
    pthread_t tids[WORKERS];
    for (long i = 0; i < WORKERS; i++) {
        pthread_create(&tids[i], NULL, worker, (void*)i);
    }

    struct epoll_event events[MAX_EVENTS];
    while (1) {
        int n = epoll_wait(epfd, events, MAX_EVENTS, -1);
        if (n < 0) { if (errno == EINTR) continue; perror("epoll_wait"); break; }

        for (int i = 0; i < n; i++) {
            int fd = events[i].data.fd;
            uint32_t evs = events[i].events;

            if (fd == lfd) {
                while (1) {
                    int conn = accept(lfd, NULL, NULL);
                    if (conn < 0) { break; }
                    set_nonblock(conn);
                    struct epoll_event cev = {
                        .events = EPOLLIN | EPOLLET | EPOLLONESHOT | EPOLLRDHUP,
                        .data.fd = conn,
                    };
                    epoll_ctl(epfd, EPOLL_CTL_ADD, conn, &cev);
                    printf("[main] 新连接 fd=%d (已加 ONESHOT)\n", conn);
                }
            } else {
                /* 派给 worker 线程处理 */
                if (evs & (EPOLLERR | EPOLLHUP | EPOLLRDHUP)) {
                    close(fd);
                    continue;
                }
                if (evs & EPOLLIN) {
                    q_push(fd);
                }
            }
        }
    }
    return 0;
}