// 练习 4：给每个连接加输出缓冲区，解决 write 写不完的问题
#include <arpa/inet.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <unistd.h>

#include <cerrno>
#include <csignal>
#include <cstdio>
#include <map>
#include <string>

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

static int epfd;
static std::map<int, std::string> outbuf;  // 每个连接还没发出去的数据

// 修改一个连接关心的事件：要不要关心 EPOLLOUT
static void watch_write(int fd, bool on) {
    epoll_event ev{};
    ev.events = on ? (EPOLLIN | EPOLLOUT) : EPOLLIN;
    ev.data.fd = fd;
    epoll_ctl(epfd, EPOLL_CTL_MOD, fd, &ev);
}

static void close_conn(int fd) {
    outbuf.erase(fd);
    epoll_ctl(epfd, EPOLL_CTL_DEL, fd, nullptr);
    close(fd);
}

// 尽量把输出缓冲区里的数据写出去；写不完的留在缓冲区里。出错返回 false
static bool flush(int fd) {
    std::string& out = outbuf[fd];
    while (!out.empty()) {
        ssize_t w = write(fd, out.data(), out.size());
        if (w > 0) {
            out.erase(0, w);  // 删掉已经写出去的部分
        } else if (w < 0 && errno == EAGAIN) {
            break;            // 内核发送缓冲区满了，等 EPOLLOUT 再写
        } else {
            return false;     // 真的出错了（比如对方已经关闭）
        }
    }
    return true;
}

int main() {
    signal(SIGPIPE, SIG_IGN);  // 往已关闭的连接写数据时，不要让 SIGPIPE 信号杀死进程

    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    int opt = 1;
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons(8888);
    if (bind(listen_fd, (sockaddr*)&addr, sizeof(addr)) < 0) { perror("bind"); return 1; }
    listen(listen_fd, 1024);
    set_nonblock(listen_fd);

    epfd = epoll_create1(0);
    epoll_event ev{};
    ev.events = EPOLLIN;
    ev.data.fd = listen_fd;
    epoll_ctl(epfd, EPOLL_CTL_ADD, listen_fd, &ev);
    printf("listening on 8888 (buffered)\n");

    epoll_event events[1024];
    char buf[65536];
    while (true) {
        int n = epoll_wait(epfd, events, 1024, -1);
        for (int i = 0; i < n; ++i) {
            int fd = events[i].data.fd;
            if (fd == listen_fd) {
                while (true) {
                    int conn_fd = accept(listen_fd, nullptr, nullptr);
                    if (conn_fd < 0) break;
                    set_nonblock(conn_fd);
                    epoll_event cev{};
                    cev.events = EPOLLIN;
                    cev.data.fd = conn_fd;
                    epoll_ctl(epfd, EPOLL_CTL_ADD, conn_fd, &cev);
                    outbuf[conn_fd];  // 创建一个空的输出缓冲区
                }
                continue;
            }

            uint32_t e = events[i].events;
            if (e & EPOLLIN) {
                ssize_t r = read(fd, buf, sizeof(buf));
                if (r > 0) {
                    outbuf[fd].append(buf, r);  // 先放进输出缓冲区
                    if (!flush(fd)) { close_conn(fd); continue; }
                    if (!outbuf[fd].empty()) watch_write(fd, true);  // 没写完，关心可写事件
                } else if (r == 0 || errno != EAGAIN) {
                    close_conn(fd);
                    continue;
                }
            }
            if (e & EPOLLOUT) {
                // 内核发送缓冲区又有空间了，接着写
                if (!flush(fd)) { close_conn(fd); continue; }
                if (outbuf[fd].empty()) watch_write(fd, false);  // 写完了，不再关心可写
            }
        }
    }
}