// What WSAPoll reports for a set that contains a closed socket.
// Winsock only, no DPP. Loopback sockets, no network.
// Build: cl /nologo /EHsc /W3 repro_wsapoll.cpp ws2_32.lib
#include <winsock2.h>
#include <chrono>
#include <cstdio>

// A connected TCP pair over 127.0.0.1.
static void tcp_pair(SOCKET &mine, SOCKET &peer) {
    SOCKET listener = socket(AF_INET, SOCK_STREAM, 0);
    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    bind(listener, (sockaddr *)&addr, sizeof addr);
    listen(listener, 1);
    int len = sizeof addr;
    getsockname(listener, (sockaddr *)&addr, &len);
    mine = socket(AF_INET, SOCK_STREAM, 0);
    connect(mine, (sockaddr *)&addr, sizeof addr);
    peer = accept(listener, nullptr, nullptr);
    closesocket(listener);
}

static void print_entry(int index, const char *what, const WSAPOLLFD &e) {
    printf("  entry %d, %-14s revents=0x%04x%s%s\n", index, what, e.revents,
           (e.revents & POLLRDNORM) ? " POLLRDNORM" : "",
           (e.revents & POLLNVAL) ? " POLLNVAL" : "");
}

int main() {
    WSADATA wsa;
    WSAStartup(MAKEWORD(2, 2), &wsa);

    SOCKET a, a_peer, b, b_peer;
    tcp_pair(a, a_peer);
    tcp_pair(b, b_peer);
    send(a_peer, "x", 1, 0);  // a and b have data to read
    send(b_peer, "x", 1, 0);
    Sleep(100);

    SOCKET closed = socket(AF_INET, SOCK_STREAM, 0);
    closesocket(closed);  // the number no longer names a socket

    printf("Test 1: [ready socket, closed socket, ready socket], timeout 1000 ms\n");
    WSAPOLLFD set[3] = {{a, POLLIN, 0}, {closed, POLLIN | POLLOUT, 0}, {b, POLLIN, 0}};
    int ret = WSAPoll(set, 3, 1000);
    printf("  WSAPoll returned %d\n", ret);
    print_entry(0, "ready socket:", set[0]);
    print_entry(1, "closed socket:", set[1]);
    print_entry(2, "ready socket:", set[2]);
    int nonzero = 0;
    for (const WSAPOLLFD &e : set) {
        nonzero += e.revents != 0;
    }
    printf("  entries with revents != 0: %d\n", nonzero);

    // The loop condition of src/dpp/socketengines/poll.cpp:75-81 applied to this result.
    printf("  DPP's loop visits entries:");
    int processed = 0;
    for (int index = 0; index < 3 && processed < ret; index++) {
        printf(" %d", index);
        if (set[index].revents > 0) {
            processed++;
        }
    }
    printf("\n");

    printf("Test 2: [closed socket, idle socket], timeout 1000 ms\n");
    WSAPOLLFD idle_set[2] = {{closed, POLLIN | POLLOUT, 0}, {a_peer, POLLIN, 0}};
    auto start = std::chrono::steady_clock::now();
    ret = WSAPoll(idle_set, 2, 1000);
    auto took = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start);
    printf("  WSAPoll returned %d after %lld ms\n", ret, (long long)took.count());
    print_entry(0, "closed socket:", idle_set[0]);
    print_entry(1, "idle socket:", idle_set[1]);
    return 0;
}

// DPP 10.1.6 on Windows: a closed socket whose entry stays in the poll engine
// makes process_events() skip the last ready socket, and the loop spins.
// Loopback sockets only: no bot token, no Discord, no network.
//
// Socket A is closed inside its own read handler, the way ssl_connection
// does it when handle_buffer() closes the connection and still returns true
// (line numbers at tag v10.1.6):
//   close():   socketengine->delete_socket(sfd); close_socket(sfd);  src/dpp/sslconnection.cpp:690-694
//   on_read(): se.flags = WANT_READ | WANT_WRITE | WANT_ERROR;
//              socketengine->update_socket(se);                       src/dpp/sslconnection.cpp:394-401
// update_socket() overwrites the whole entry, WANT_DELETION included
// (src/dpp/socketengine.cpp:55), so the engine never removes it
// (src/dpp/socketengines/poll.cpp:131-135).
// The control run closes A the same way but without the update_socket() call.
#include <dpp/dpp.h>
#include <winsock2.h>
#include <atomic>
#include <chrono>
#include <cstdio>

using namespace std::chrono;

// A connected TCP pair over 127.0.0.1.
static void tcp_pair(SOCKET &mine, SOCKET &peer) {
    SOCKET listener = socket(AF_INET, SOCK_STREAM, 0);
    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    bind(listener, (sockaddr *)&addr, sizeof addr);
    listen(listener, 1);
    int len = sizeof addr;
    getsockname(listener, (sockaddr *)&addr, &len);
    mine = socket(AF_INET, SOCK_STREAM, 0);
    connect(mine, (sockaddr *)&addr, sizeof addr);
    peer = accept(listener, nullptr, nullptr);
    closesocket(listener);
}

static dpp::socket_events reader(SOCKET s, int &calls) {
    return dpp::socket_events(s, dpp::WANT_READ | dpp::WANT_ERROR,
        [&calls](dpp::socket fd, const dpp::socket_events &) {
            char buf[16];
            recv(fd, buf, sizeof buf, 0);
            ++calls;
        },
        {}, [](dpp::socket, const dpp::socket_events &, int) {});
}

// Calls process_events() for one second. A call that is waiting in WSAPoll when
// the second ends is finished and counted.
static long run_one_second(dpp::socket_engine_base &engine) {
    long passes = 0;
    auto end = steady_clock::now() + seconds(1);
    while (steady_clock::now() < end) {
        engine.process_events();
        ++passes;
    }
    return passes;
}

static void run(bool rearm) {
    printf(rearm ? "A closed in its read handler, then update_socket() (as ssl_connection::on_read does):\n"
                 : "Control: A closed in its read handler, no update_socket():\n");
    int g_calls = 0, b_calls = 0, c_calls = 0;
    std::atomic<bool> a_closed{false}, a_removed{false};

    // All sockets exist before A is closed, so none of them can get A's number
    // (a socket registered with that number replaces A's entry,
    // src/dpp/socketengine.cpp:41-46).
    SOCKET g, g_peer, a, a_peer, b, b_peer, c, c_peer;
    tcp_pair(g, g_peer);
    tcp_pair(a, a_peer);
    tcp_pair(b, b_peer);
    tcp_pair(c, c_peer);

    dpp::cluster bot;  // no token; start() is never called
    dpp::socket_engine_base &engine = *bot.socketengine;
    bot.on_socket_close([&](const dpp::socket_close_t &e) {
        if (e.fd == a) {
            a_removed = true;
        }
    });

    engine.register_socket(reader(g, g_calls));
    engine.register_socket(dpp::socket_events(a, dpp::WANT_READ | dpp::WANT_ERROR,
        [&](dpp::socket fd, const dpp::socket_events &ev) {
            char buf[16];
            recv(fd, buf, sizeof buf, 0);
            engine.delete_socket(fd);
            closesocket(fd);
            if (rearm) {
                dpp::socket_events se{ev};
                se.flags = dpp::WANT_READ | dpp::WANT_WRITE | dpp::WANT_ERROR;
                engine.update_socket(se);
            }
            a_closed = true;
        },
        {}, [](dpp::socket, const dpp::socket_events &, int) {}));
    send(a_peer, "x", 1, 0);
    auto deadline = steady_clock::now() + seconds(3);
    while (!a_closed && steady_clock::now() < deadline) {
        engine.process_events();
    }
    printf("  A closed: %s; engine removed A's entry: %s\n", a_closed ? "yes" : "no", a_removed ? "yes" : "no");

    engine.register_socket(reader(b, b_calls));
    send(g_peer, "x", 1, 0);
    send(b_peer, "x", 1, 0);
    long passes = run_one_second(engine);
    printf("  G (registered before A) and B (after A) get data: G read calls=%d, B read calls=%d, passes in 1 s=%ld\n",
           g_calls, b_calls, passes);

    engine.register_socket(reader(c, c_calls));
    send(c_peer, "x", 1, 0);
    passes = run_one_second(engine);
    printf("  C (after B) gets data:                            B read calls=%d, C read calls=%d, passes in 1 s=%ld\n",
           b_calls, c_calls, passes);
    printf("  engine removed A's entry by the end: %s\n\n", a_removed ? "yes" : "no");

    for (SOCKET s : {g, g_peer, a_peer, b, b_peer, c, c_peer}) {
        closesocket(s);
    }
}

int main() {
    WSADATA wsa;
    WSAStartup(MAKEWORD(2, 2), &wsa);
    run(true);
    run(false);
    return 0;
}