//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

// ClientQUIC::OnClosed tells the owner of a connection that it is over, so
// the owner can delete the ClientQUIC. A ListenerQUIC hands each accepted
// connection to its callback as a heap-allocated ClientQUIC the callback
// owns, and nothing said when to delete it: msquic's handle of a connection
// the peer closed was closed for it, but the object stayed, so a server
// either leaked its peers or deleted them blind. One still alive when main
// returns holds msquic's registration open at exit until its idle timeout.
//
// The server here deletes every connection it accepts from inside its
// OnClosed callback, which runs on msquic's worker, over 25 rounds of four
// clients that connect, echo and close; and main returns normally, so
// msquic's RegistrationClose at exit is the check that nothing was left
// open, under a 45 s watchdog.
//
// Also, each exactly once:
//  - a peer that vanishes is reported: a forked child that _Exits
//    mid-connection (its closed port answers the server's next packet, a
//    transport error), and one that stops itself with SIGSTOP, keeping its
//    port and never answering, which only the idle timeout ends;
//  - Stop() on the server side fires it there, and the client's own
//    OnClosed hears the server go;
//  - a connection that closed before OnClosed was set still gets it, from a
//    pool task, not from inside OnClosed;
//  - an owner that deletes the ClientQUIC first never hears it.

#include <sys/wait.h>
#include <unistd.h>
#include <cerrno>
#include <csignal>
#include <cstdio>

import Crafter.Network;
import Crafter.Thread;
import std;
using namespace Crafter;

namespace {

int failures = 0;

void Check(bool ok, std::string_view what) {
    std::println("{}  {}", ok ? "PASS" : "FAIL", what);
    if (!ok) ++failures;
}

const std::string alpn = "crafter/test-on-closed";

QUICClientCredentials ClientCredentials() {
    QUICClientCredentials creds;
    creds.insecureNoServerValidation = true;
    return creds;
}

QUICServerCredentials ServerCredentials(std::uint32_t idleTimeoutMs = 120'000) {
    QUICServerCredentials creds;
    creds.selfSigned = true;
    creds.settings.idleTimeoutMs = idleTimeoutMs;
    return creds;
}

void Echo(ClientQUIC* peer) {
    peer->OnStream([](QUICStream stream) {
        try {
            std::vector<char> got = stream.RecieveUntilCloseSync();
            stream.SendSync(got.data(), static_cast<std::uint32_t>(got.size()), true);
        } catch (...) {}
    });
}

bool RoundTrip(ClientQUIC& client, std::string_view text) {
    QUICStream stream = client.OpenStream();
    stream.SendSync(text.data(), static_cast<std::uint32_t>(text.size()), true);
    std::vector<char> echoed = stream.RecieveUntilCloseSync();
    return std::string(echoed.begin(), echoed.end()) == text;
}

// Waits, up to `limit`, for `done` to hold.
bool WaitFor(std::function<bool()> done, std::chrono::milliseconds limit) {
    const auto deadline = std::chrono::steady_clock::now() + limit;
    while (!done()) {
        if (std::chrono::steady_clock::now() >= deadline) return false;
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
    }
    return true;
}

// The server: owns what it accepted, and deletes it when told it closed.
struct Server {
    std::mutex mutex;
    std::unordered_set<ClientQUIC*> live;
    std::vector<std::shared_ptr<std::atomic<int>>> fired;   // one per connection
    std::atomic<int> accepted = 0;
    std::atomic<int> notified = 0;

    // Takes ownership of `peer`; `echo` serves its streams, as a server
    // would, which the Stop() case replaces with a handler of its own.
    void Adopt(ClientQUIC* peer, bool echo = true) {
        auto count = std::make_shared<std::atomic<int>>(0);
        {
            std::lock_guard lock(mutex);
            live.insert(peer);
            fired.push_back(count);
        }
        ++accepted;
        peer->OnClosed([this, peer, count] {
            ++*count;
            ++notified;
            {
                std::lock_guard lock(mutex);
                live.erase(peer);
            }
            // From inside the notification, on msquic's worker.
            delete peer;
        });
        if (echo) Echo(peer);
    }

    bool AllFiredOnce() {
        std::lock_guard lock(mutex);
        return std::ranges::all_of(fired, [](const auto& count) { return count->load() == 1; });
    }
};

// A vanishing peer, in a process of its own: connect, prove the
// connection is up with one echo, and disappear without a CONNECTION_CLOSE,
// by exiting or by freezing.
[[noreturn]] void Vanish(int portPipe, bool freeze) {
    std::uint16_t port = 0;
    while (read(portPipe, &port, sizeof(port)) < 0 && errno == EINTR) {}
    try {
        ThreadPool::Start();
        ClientQUIC client("127.0.0.1", port, alpn, ClientCredentials());
        (void)RoundTrip(client, "going");
        if (freeze) raise(SIGSTOP);
    } catch (...) {}
    std::_Exit(0);
}

struct Vanisher {
    pid_t pid = -1;
    int portPipe = -1;
};

// Forks a vanishing peer, waiting for the port of its server on a pipe.
Vanisher ForkVanisher(bool freeze) {
    int ends[2];
    if (pipe(ends) != 0) return {};
    const pid_t pid = fork();
    if (pid == 0) {
        close(ends[1]);
        Vanish(ends[0], freeze);
    }
    close(ends[0]);
    return { pid, ends[1] };
}

} // namespace

int main() {
    // Before any thread exists, so the children start clean.
    const Vanisher exits = ForkVanisher(false);
    const Vanisher freezes = ForkVanisher(true);
    if (exits.pid < 0 || freezes.pid < 0) return 1;

    ThreadPool::Start();
    std::thread([] {
        std::this_thread::sleep_for(std::chrono::seconds(45));
        std::println("FAIL  timed out: a notification never came, or the exit hung");
        std::fflush(stdout);
        std::_Exit(1);
    }).detach();

    try {
        // ── Delete from inside the notification, over 100 connections ─────────
        // Each Server is leaked on purpose: a notification that came after
        // its block on a failure path must not find it destroyed.
        Server& server = *new Server;
        {
            ListenerQUIC listener("127.0.0.1", 0, alpn, ServerCredentials(),
                                  [&](ClientQUIC* peer) { server.Adopt(peer); });
            listener.ListenAsyncAsync();
            std::atomic<int> echoed = 0;
            for (int round = 0; round < 25; ++round) {
                std::vector<std::thread> clients;
                for (int i = 0; i < 4; ++i) {
                    clients.emplace_back([&] {
                        try {
                            ClientQUIC client("127.0.0.1", listener.Port(), alpn, ClientCredentials());
                            if (RoundTrip(client, "ping")) ++echoed;
                        } catch (const std::exception& e) {
                            std::println("      client: {}", e.what());
                        }
                        // ~ClientQUIC closes the connection: the peer close
                        // the server is told about.
                    });
                }
                for (std::thread& client : clients) client.join();
            }
            Check(echoed.load() == 100, std::format("100 connections served ({})", echoed.load()));
            const bool all = WaitFor([&] { return server.notified.load() == server.accepted.load(); },
                                     std::chrono::seconds(10));
            Check(all && server.accepted.load() == 100,
                  std::format("every connection a peer closed was reported closed ({} of {})",
                              server.notified.load(), server.accepted.load()));
            Check(server.AllFiredOnce(), "each exactly once");
            std::lock_guard lock(server.mutex);
            Check(server.live.empty(), "and deleting each from inside its notification left none behind");
        }

        // ── Peers that vanish ────────────────────────────────────────────────
        // Each against a listener whose idle timeout is 1.5 s; `stopped` is
        // when waitpid saw the child exit or stop.
        auto vanish = [](const Vanisher& child, int waitFlags, std::string_view how) {
            Server& server = *new Server;
            ListenerQUIC listener("127.0.0.1", 0, alpn, ServerCredentials(1500),
                                  [&server](ClientQUIC* peer) { server.Adopt(peer); });
            listener.ListenAsyncAsync();
            const std::uint16_t port = listener.Port();
            (void)!write(child.portPipe, &port, sizeof(port));
            close(child.portPipe);
            int status = 0;
            while (waitpid(child.pid, &status, waitFlags) < 0 && errno == EINTR) {}
            const auto stopped = std::chrono::steady_clock::now();
            const bool notified = WaitFor([&] { return server.notified.load() == 1; }, std::chrono::seconds(10));
            const auto after = std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - stopped);
            Check(server.accepted.load() == 1 && notified && server.AllFiredOnce(),
                  std::format("a peer that {} is reported closed, once ({} ms later)", how, after.count()));
            return after;
        };
        vanish(exits, 0, "exited without closing");
        const std::chrono::milliseconds frozen = vanish(freezes, WUNTRACED, "froze, keeping its port,");
        Check(frozen >= std::chrono::milliseconds(1000),
              "the frozen one only once the idle timeout ran out");
        kill(freezes.pid, SIGKILL);
        while (waitpid(freezes.pid, nullptr, 0) < 0 && errno == EINTR) {}

        // ── Stop() on the server side; the client hears it too ───────────────
        {
            Server& stopping = *new Server;
            // Stop once the client's first stream arrives, so the handshake
            // is over: stopping in connectCallback can beat the handshake,
            // and the client's constructor then throws instead.
            ListenerQUIC listener("127.0.0.1", 0, alpn, ServerCredentials(), [&](ClientQUIC* peer) {
                stopping.Adopt(peer, false);
                peer->OnStream([peer](QUICStream) { peer->Stop(); });
            });
            listener.ListenAsyncAsync();
            std::atomic<int> clientHeard = 0;
            {
                ClientQUIC client("127.0.0.1", listener.Port(), alpn, ClientCredentials());
                client.OnClosed([&] { ++clientHeard; });
                QUICStream stream = client.OpenStream();
                stream.SendSync("stop", 4, false);
                Check(WaitFor([&] { return stopping.notified.load() == 1; }, std::chrono::seconds(10))
                          && stopping.AllFiredOnce(),
                      "Stop() on the server side fires the server's notification, once");
                Check(WaitFor([&] { return clientHeard.load() == 1; }, std::chrono::seconds(10)),
                      "and the client's own OnClosed hears the server close it");
            }
            Check(clientHeard.load() == 1, "destroying the client after that does not fire it again");
        }

        // ── Closed before OnClosed was set ───────────────────────────────────
        {
            // Leaked, as the Servers are: connectCallback may still be
            // returning from OnClosed when the check below has passed.
            struct Late {
                std::atomic<bool> clientGone = false;
                std::atomic<bool> inside = false;
                std::atomic<int> fired = 0;
                std::atomic<bool> firedInside = false;
            };
            Late& late = *new Late;
            ListenerQUIC listener("127.0.0.1", 0, alpn, ServerCredentials(), [&late](ClientQUIC* peer) {
                // Long after the peer closed: the close has been and gone.
                WaitFor([&] { return late.clientGone.load(); }, std::chrono::seconds(10));
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
                late.inside = true;
                peer->OnClosed([&late, peer] {
                    ++late.fired;
                    late.firedInside = late.inside.load();
                    delete peer;
                });
                late.inside = false;
            });
            listener.ListenAsyncAsync();
            {
                ClientQUIC client("127.0.0.1", listener.Port(), alpn, ClientCredentials());
            }
            late.clientGone = true;
            Check(WaitFor([&] { return late.fired.load() == 1; }, std::chrono::seconds(10)),
                  "a connection that closed before OnClosed was set is still reported");
            Check(!late.firedInside.load(), "from a pool task, not from inside OnClosed");
        }

        // ── The owner deletes first ──────────────────────────────────────────
        {
            struct Counts {
                std::atomic<int> heard = 0;
                std::atomic<int> deleted = 0;
            };
            Counts& counts = *new Counts;
            std::atomic<int>& heard = counts.heard;
            std::atomic<int>& deleted = counts.deleted;
            ListenerQUIC listener("127.0.0.1", 0, alpn, ServerCredentials(), [&counts](ClientQUIC* peer) {
                peer->OnClosed([&counts] { ++counts.heard; });
                delete peer;
                ++counts.deleted;
            });
            listener.ListenAsyncAsync();
            for (int i = 0; i < 5; ++i) {
                try {
                    ClientQUIC client("127.0.0.1", listener.Port(), alpn, ClientCredentials());
                } catch (...) {}
            }
            WaitFor([&] { return deleted.load() == 5; }, std::chrono::seconds(10));
            std::this_thread::sleep_for(std::chrono::milliseconds(300));
            Check(deleted.load() == 5 && heard.load() == 0,
                  std::format("an owner that deletes the connection first never hears it ({} heard)", heard.load()));
        }
    } catch (const std::exception& e) {
        Check(false, std::format("threw: {}", e.what()));
    }

    std::println("{}", failures == 0 ? "all checks passed" : "check(s) failed");
    std::fflush(stdout);
    // Deliberately no std::_Exit: returning runs msquic's static
    // RegistrationClose, which blocks until every connection is closed. Every
    // accepted ClientQUIC was deleted by its notification or its owner, so
    // this returns at once; a leaked one would hold it until the watchdog.
    return failures == 0 ? 0 : 1;
}