#include <unistd.h>
#include <sys/epoll.h>
#include <curl/curl.h>
#include <vector>
#include <map>
#include <iostream>
#include <set>

struct pollset {
public:
    pollset() {
        epollfd = epoll_create1(0);
        timeout_ms = 0;
        multi = curl_multi_init();
        curl_multi_setopt(multi, CURLMOPT_SOCKETFUNCTION, &pollset::socket_callback);
        curl_multi_setopt(multi, CURLMOPT_SOCKETDATA, this);

        curl_multi_setopt(multi, CURLMOPT_TIMERFUNCTION, &pollset::timer_callback);
        curl_multi_setopt(multi, CURLMOPT_TIMERDATA, this);
    }

    ~pollset() {
        curl_multi_cleanup(multi);
        close(epollfd);
    }

    void add(CURL* handle) {
        curl_multi_add_handle(multi, handle);
    }

    void remove(CURL* handle) {
        curl_multi_remove_handle(multi, handle);
    }

    void set_event(int socket, uint32_t events) {
        auto inserted = handled_fds.insert(socket);
        auto op = (inserted.second)?EPOLL_CTL_ADD:EPOLL_CTL_MOD;
        epoll_event e;
        e.events = events;
        e.data.fd = socket;
        epoll_ctl(epollfd, op, socket, &e);
    }

    void remove_event(int socket) {
        handled_fds.erase(socket);
        epoll_ctl(epollfd, EPOLL_CTL_DEL, socket, nullptr);
    }

    void set_timeout(int timeout_ms) {
        this->timeout_ms = timeout_ms;
    }

    CURLMsg *info_read(int& msgs_in_queue) {
        return curl_multi_info_read(multi, &msgs_in_queue);
    }

    void poll() {
        epoll_event evt;
        auto status = epoll_wait(epollfd, &evt, 1, timeout_ms);

        if (status == 0) {
            perform_action(CURL_SOCKET_TIMEOUT, 0);
        } else {
            int action = 0;
            if (evt.events & EPOLLIN)
                action |= CURL_CSELECT_IN;
            if (evt.events & EPOLLOUT)
                action |= CURL_CSELECT_OUT;
            if (evt.events & EPOLLERR)
                action |= CURL_CSELECT_ERR;
            perform_action(evt.data.fd, action);
        }
    }

private:
    int perform_action(int fd, int action) {
            int running;
            while (true) {
                auto ret = curl_multi_socket_action(multi, fd, action,
                                                    &running);
                if (ret != CURLM_CALL_MULTI_PERFORM)
                    break ;
            }
            return running;
        }

private:
    static int socket_callback(CURL *easy,
                               curl_socket_t s,
                               int what,
                               void* userp,
                               void* /*socketp*/) {
        auto polls = static_cast<pollset*>(userp);
        switch (what) {
        case CURL_POLL_IN:
            polls->set_event(s, EPOLLIN);
            break ;
        case CURL_POLL_OUT:
            polls->set_event(s, EPOLLOUT);
            break ;
        case CURL_POLL_INOUT:
            polls->set_event(s, EPOLLIN | EPOLLOUT);
            break ;
        case CURL_POLL_REMOVE:
            polls->remove_event(s);
        default:
            break ;
        }

        return 0;
    }

    static int timer_callback(CURLM *multi,
                              long timeout_ms,
                              void *userp) {
        auto polls = static_cast<pollset*>(userp);
        polls->set_timeout(timeout_ms);
    }

    CURLM* multi;

    int timeout_ms;
    std::set<int> handled_fds;
    int epollfd;
};

int main(void)
{
    pollset p;

    std::vector<CURL*> handles;
    for (unsigned i = 0; i < 100; i++) {
        handles.push_back(curl_easy_init());
        curl_easy_setopt(handles[i], CURLOPT_URL, "http://localhost:5432/");
        curl_easy_setopt(handles[i], CURLOPT_VERBOSE, 1L);
        p.add(handles[i]);
    }

    auto requests = handles.size();

    while (requests > 0) {
        p.poll();
        while (true) {
            int msgq = 0;
            auto m = p.info_read(msgq);
            if (!m)
                break ;
            if(m->msg == CURLMSG_DONE) {
                --requests;
            }
        }
    }

    for (auto& h : handles) {
        p.remove(h);
        curl_easy_cleanup(h);
    }

    return 0;
}