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

// QUICStream::SendAsync holds no thread. It used to be SendSync run as a
// pool task: a pool thread was held until msquic accepted the buffer, which
// a slow peer stretches to its own pace, so a server streaming to a slow
// follower lost a worker of a pool the host shares, and the task reached back
// through the wrapper, which a move or a destructor in the meantime freed.
// Now the copy goes to msquic at once and its SEND_COMPLETE reports, on
// msquic's worker, exactly once.
//
//  - A burst of 256 sends of 64 KiB returns with most of it still in flight
//    (BytesInFlight), every send then reports true, in order, the counters
//    drain to zero, and the peer gets every byte in order.
//  - With every pool worker parked, a chain of 100 sends, each made from the
//    previous one's onSent, completes: nothing needed the pool.
//  - A connection stopped under a burst reports the rest false, each send
//    exactly once, the trues before the falses.
//  - Sends made just before the wrapper is destroyed still go out and
//    report true.
//  - Four threads sending on one stream at once all complete.
//  - A send on a closed stream reports false from a pool task, not from
//    inside SendAsync.

#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;
}

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(2));
    }
    return true;
}

// Byte `offset` of a stream, after its one-byte tag.
char Pattern(std::uint64_t offset) { return static_cast<char>((offset * 131 + 7) & 0xff); }

// What the server read from each stream, by the stream's first byte.
struct Received {
    std::mutex mutex;
    std::map<char, std::vector<char>> streams;

    std::optional<std::vector<char>> Of(char tag) {
        std::lock_guard lock(mutex);
        auto it = streams.find(tag);
        if (it == streams.end()) return std::nullopt;
        return it->second;
    }
};

bool FollowsPattern(const std::vector<char>& bytes, std::size_t expected) {
    if (bytes.size() != expected + 1) return false;
    for (std::size_t i = 1; i < bytes.size(); ++i) {
        if (bytes[i] != Pattern(i - 1)) return false;
    }
    return true;
}

// The completions of a set of sends, by index.
struct Completions {
    std::mutex mutex;
    std::vector<int> calls;        // how often each send's onSent fired
    std::vector<bool> sent;        // what it said
    std::vector<int> order;        // indices, in the order they fired

    explicit Completions(std::size_t count) : calls(count, 0), sent(count, false) {}

    std::function<void(bool)> For(std::size_t index) {
        return [this, index](bool ok) {
            std::lock_guard lock(mutex);
            ++calls[index];
            sent[index] = ok;
            order.push_back(static_cast<int>(index));
        };
    }
    std::size_t Fired() {
        std::lock_guard lock(mutex);
        return order.size();
    }
};

} // namespace

int main() {
    ThreadPool::Start();
    std::thread([] {
        std::this_thread::sleep_for(std::chrono::seconds(45));
        std::println("FAIL  timed out");
        std::fflush(stdout);
        std::_Exit(1);
    }).detach();

    const std::string alpn = "crafter/test-send-async";
    QUICServerCredentials serverCreds;
    serverCreds.selfSigned = true;
    QUICClientCredentials clientCreds;
    clientCreds.insecureNoServerValidation = true;

    // Leaked: readers may still be finishing when main has moved on.
    Received& received = *new Received;
    std::atomic<int> readersStarted = 0;

    try {
        // Each stream gets a reader on a thread of its own, not on the pool,
        // which the second part parks.
        ListenerQUIC listener("127.0.0.1", 0, alpn, serverCreds, [&](ClientQUIC* peer) {
            peer->OnStream([&](QUICStream stream) {
                auto shared = std::make_shared<QUICStream>(std::move(stream));
                ++readersStarted;
                std::thread([&received, shared] {
                    std::vector<char> bytes;
                    try {
                        bytes = shared->RecieveUntilCloseSync();
                    } catch (...) {}
                    if (bytes.empty()) return;
                    std::lock_guard lock(received.mutex);
                    received.streams[bytes.front()] = std::move(bytes);
                }).detach();
            });
        });
        listener.ListenAsyncAsync();
        ClientQUIC client("127.0.0.1", listener.Port(), alpn, clientCreds);

        // ── A burst ──────────────────────────────────────────────────────────
        {
            constexpr std::size_t count = 256;
            constexpr std::size_t size = 64 * 1024;
            auto completions = std::make_shared<Completions>(count + 1);
            QUICStream stream = client.OpenStream();
            stream.SendAsync("A", 1, false, completions->For(count));
            std::vector<char> chunk(size);
            const auto start = std::chrono::steady_clock::now();
            for (std::size_t i = 0; i < count; ++i) {
                for (std::size_t j = 0; j < size; ++j) chunk[j] = Pattern(i * size + j);
                stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(size), i + 1 == count, completions->For(i));
            }
            const auto took = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start);
            const std::uint64_t inFlight = stream.BytesInFlight();
            Check(inFlight > 0 && stream.SendsInFlight() > 0,
                  std::format("256 sends of 64 KiB returned in {} ms with {} KiB still in flight: none waited for msquic",
                              took.count(), inFlight / 1024));
            Check(WaitFor([&] { return completions->Fired() == count + 1; }, std::chrono::seconds(20)),
                  "every send reported");
            {
                std::lock_guard lock(completions->mutex);
                const bool once = std::ranges::all_of(completions->calls, [](int calls) { return calls == 1; });
                const bool allSent = std::ranges::all_of(completions->sent, [](bool ok) { return ok; });
                std::vector<int> data(completions->order.begin(), completions->order.end());
                std::erase(data, static_cast<int>(count));
                Check(once && allSent && std::ranges::is_sorted(data), "each exactly once, true, and in the order sent");
            }
            Check(stream.BytesInFlight() == 0 && stream.SendsInFlight() == 0, "after which nothing is in flight");
            const bool arrived = WaitFor([&] { return received.Of('A').has_value(); }, std::chrono::seconds(20));
            Check(arrived && FollowsPattern(*received.Of('A'), count * size), "and the peer read all 16 MiB, in order");
        }

        // ── No pool ──────────────────────────────────────────────────────────
        {
            QUICStream stream = client.OpenStream();
            const int before = readersStarted.load();
            stream.SendSync("B", 1, false);
            // The server's OnStream is itself a pool task: let it start the
            // reader before the pool is parked.
            WaitFor([&] { return readersStarted.load() > before; }, std::chrono::seconds(10));

            // Leaked, and released only once every parked task has left, so
            // none reads it after this block.
            struct Parking {
                std::atomic<bool> release = false;
                std::atomic<unsigned> parked = 0;
                std::atomic<unsigned> left = 0;
            };
            Parking& parking = *new Parking;
            const unsigned workers = ThreadPool::activeThreads.load();
            for (unsigned i = 0; i < workers * 2; ++i) {
                ThreadPool::Enqueue([&parking] {
                    ++parking.parked;
                    while (!parking.release.load()) std::this_thread::sleep_for(std::chrono::milliseconds(1));
                    ++parking.left;
                });
            }
            const bool allParked = WaitFor([&] { return parking.parked.load() >= workers; }, std::chrono::seconds(10));

            constexpr int links = 100;
            auto done = std::make_shared<std::atomic<int>>(0);
            auto ok = std::make_shared<std::atomic<bool>>(true);
            // Each onSent makes the next send, on msquic's worker.
            auto chain = std::make_shared<std::function<void(int)>>();
            QUICStream* raw = &stream;
            *chain = [raw, chain, done, ok](int link) {
                const char byte = Pattern(static_cast<std::uint64_t>(link));
                raw->SendAsync(&byte, 1, link + 1 == links, [raw, chain, done, ok, link](bool sent) {
                    if (!sent) *ok = false;
                    ++*done;
                    if (link + 1 < links) (*chain)(link + 1);
                });
            };
            (*chain)(0);
            const bool finished = WaitFor([&] { return done->load() == links; }, std::chrono::seconds(10));
            Check(allParked && finished && ok->load(),
                  std::format("with all {} pool workers parked, a chain of {} sends, each from the last one's onSent, completed",
                              workers, done->load()));
            parking.release = true;
            WaitFor([&] { return parking.left.load() == workers * 2; }, std::chrono::seconds(10));
            const bool arrived = WaitFor([&] { return received.Of('B').has_value(); }, std::chrono::seconds(10));
            Check(arrived && FollowsPattern(*received.Of('B'), links), "and the peer read them all, in order");
            *chain = nullptr;
        }

        // ── The connection stopped under a burst ─────────────────────────────
        {
            ClientQUIC doomed("127.0.0.1", listener.Port(), alpn, clientCreds);
            QUICStream stream = doomed.OpenStream();
            constexpr std::size_t count = 256;
            constexpr std::size_t size = 64 * 1024;
            auto completions = std::make_shared<Completions>(count);
            std::vector<char> chunk(size, 'C');
            for (std::size_t i = 0; i < count; ++i) {
                stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(size), false, completions->For(i));
            }
            doomed.Stop();
            Check(WaitFor([&] { return completions->Fired() == count; }, std::chrono::seconds(20)),
                  "a connection stopped under a burst reports every send");
            std::lock_guard lock(completions->mutex);
            const bool once = std::ranges::all_of(completions->calls, [](int calls) { return calls == 1; });
            const auto firstFalse = std::ranges::find(completions->sent, false);
            const bool prefix = std::all_of(firstFalse, completions->sent.end(), [](bool ok) { return !ok; });
            Check(once && prefix && firstFalse != completions->sent.end(),
                  std::format("each exactly once: {} true, then {} false",
                              firstFalse - completions->sent.begin(), completions->sent.end() - firstFalse));
            Check(stream.BytesInFlight() == 0, "and nothing is left in flight");
        }

        // ── The wrapper destroyed with sends in flight ───────────────────────
        {
            constexpr std::size_t count = 64;
            constexpr std::size_t size = 64 * 1024;
            auto completions = std::make_shared<Completions>(count + 1);
            {
                QUICStream stream = client.OpenStream();
                stream.SendAsync("D", 1, false, completions->For(count));
                std::vector<char> chunk(size);
                for (std::size_t i = 0; i < count; ++i) {
                    for (std::size_t j = 0; j < size; ++j) chunk[j] = Pattern(i * size + j);
                    stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(size), false, completions->For(i));
                }
                // ~QUICStream: Stop(), a FIN after what was sent.
            }
            Check(WaitFor([&] { return completions->Fired() == count + 1; }, std::chrono::seconds(20)),
                  "sends made just before the wrapper was destroyed all report");
            {
                std::lock_guard lock(completions->mutex);
                Check(std::ranges::all_of(completions->sent, [](bool ok) { return ok; }), "and report true");
            }
            const bool arrived = WaitFor([&] { return received.Of('D').has_value(); }, std::chrono::seconds(20));
            Check(arrived && FollowsPattern(*received.Of('D'), count * size), "and go out, followed by the FIN");
        }

        // ── Four threads on one stream ───────────────────────────────────────
        {
            constexpr int threads = 4;
            constexpr int each = 50;
            auto completions = std::make_shared<Completions>(threads * each + 1);
            QUICStream stream = client.OpenStream();
            stream.SendAsync("F", 1, false, completions->For(threads * each));
            std::vector<std::thread> senders;
            for (int t = 0; t < threads; ++t) {
                senders.emplace_back([&, t] {
                    const std::vector<char> message(1024, static_cast<char>('a' + t));
                    for (int i = 0; i < each; ++i) {
                        stream.SendAsync(message.data(), static_cast<std::uint32_t>(message.size()), false,
                                         completions->For(static_cast<std::size_t>(t * each + i)));
                    }
                });
            }
            for (std::thread& sender : senders) sender.join();
            stream.SendAsync(nullptr, 0, true, nullptr);
            Check(WaitFor([&] { return completions->Fired() == threads * each + 1; }, std::chrono::seconds(20)),
                  "four threads sending on one stream at once all hear back");
            const bool arrived = WaitFor([&] { return received.Of('F').has_value(); }, std::chrono::seconds(20));
            Check(arrived && received.Of('F')->size() == 1 + threads * each * 1024, "and the peer read every byte");
        }

        // ── A closed stream ──────────────────────────────────────────────────
        {
            QUICStream closed;
            auto inside = std::make_shared<std::atomic<bool>>(false);
            auto fired = std::make_shared<std::atomic<int>>(0);
            auto sentInside = std::make_shared<std::atomic<bool>>(false);
            auto said = std::make_shared<std::atomic<bool>>(true);
            *inside = true;
            closed.SendAsync("x", 1, false, [inside, fired, sentInside, said](bool ok) {
                *sentInside = inside->load();
                *said = ok;
                ++*fired;
            });
            *inside = false;
            Check(WaitFor([&] { return fired->load() == 1; }, std::chrono::seconds(10)) && !said->load() && !sentInside->load(),
                  "a send on a closed stream reports false, from a pool task rather than inside SendAsync");
        }
    } catch (const std::exception& e) {
        Check(false, std::format("threw: {}", e.what()));
    }

    std::println("{}", failures == 0 ? "all checks passed" : "check(s) failed");
    // As the other QUIC tests: skip the static teardown, which waits for the
    // accepted peer's idle timeout. _Exit doesn't flush.
    std::fflush(stdout);
    std::_Exit(failures == 0 ? 0 : 1);
}