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

// QUICStream::OnRecieve delivers a stream's bytes and its end by callback,
// holding no thread while nothing arrives. Until now every way to wait for
// bytes held one: RecieveSync and RecieveUntilCloseSync block their caller
// with no bound, and RecieveAsync and RecieveUntilCloseAsync block a worker
// of the pool the host shares, which for a stream a client keeps open (a log
// follower) is for good. A server embedded in an app cannot afford either.
//
//  - Bytes that arrived before OnRecieve are delivered by it, before it
//    returns, and the end with them: clean for a FIN.
//  - Registered while bytes keep arriving, with a callback slow enough to
//    let the registering thread and msquic's worker overlap, and an echo
//    sent from inside onData: no two callbacks of the stream ever run at
//    once, and every byte arrives in order.
//  - With every pool worker parked, a stream is echoed both ways, callbacks
//    and SendAsync only: nothing needed the pool.
//  - The end fires once: not clean for the peer's Abort() and for a lost
//    connection.
//  - Abort() from inside onData enforces a byte cap: no onData after it,
//    onEnd(false) once, and the peer sees both directions reset.
//  - The wrapper destroyed inside its own onEnd, or moved while bytes flow,
//    breaks nothing; destroyed from another thread mid-flow, no callback
//    runs once its destructor has returned.
//  - A second OnRecieve throws; a stream that cannot receive, or no stream,
//    gets onEnd(false) at once.

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

char Pattern(std::uint64_t offset) { return static_cast<char>((offset * 131 + 7) & 0xff); }

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

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

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

// What one stream's callbacks saw. Shared with them, and leaked, so a
// callback that comes late on a failure path finds it alive.
struct Seen {
    std::mutex mutex;
    std::string bytes;
    int ends = 0;
    bool clean = false;
    int dataAfterEnd = 0;
    std::atomic<int> inside = 0;      // callbacks of this stream running now
    std::atomic<int> overlaps = 0;    // times two ran at once

    void Enter() { if (inside.fetch_add(1) != 0) ++overlaps; }
    void Leave() { --inside; }

    std::function<void(std::span<const char>)> OnData() {
        return [this](std::span<const char> data) {
            Enter();
            {
                std::lock_guard lock(mutex);
                if (ends != 0) ++dataAfterEnd;
                bytes.append(data.begin(), data.end());
            }
            Leave();
        };
    }
    std::function<void(bool)> OnEnd() {
        return [this](bool wasClean) {
            Enter();
            {
                std::lock_guard lock(mutex);
                ++ends;
                clean = wasClean;
            }
            Leave();
        };
    }
    int Ends() {
        std::lock_guard lock(mutex);
        return ends;
    }
    std::string Bytes() {
        std::lock_guard lock(mutex);
        return bytes;
    }
};

// Every pool worker parked until Release(), which returns once each has
// left, so nothing reads this after its owner is gone.
struct Parking {
    std::atomic<bool> release = false;
    std::atomic<unsigned> parked = 0;
    std::atomic<unsigned> left = 0;
    unsigned tasks = 0;

    bool Park() {
        const unsigned workers = ThreadPool::activeThreads.load();
        tasks = workers * 2;
        for (unsigned i = 0; i < tasks; ++i) {
            ThreadPool::Enqueue([this] {
                ++parked;
                while (!release.load()) std::this_thread::sleep_for(std::chrono::milliseconds(1));
                ++left;
            });
        }
        return WaitFor([&] { return parked.load() >= workers; }, std::chrono::seconds(10));
    }
    void Release() {
        release = true;
        WaitFor([&] { return left.load() == tasks; }, std::chrono::seconds(10));
    }
};

// A listener whose every accepted stream goes to `onStream`, as ListenerQUIC
// delivers it: on a pool task. Wrappers the handler keeps go in `kept`.
struct Server {
    std::mutex mutex;
    std::vector<std::shared_ptr<QUICStream>> kept;
    std::unique_ptr<ListenerQUIC> listener;

    explicit Server(std::function<void(Server&, std::shared_ptr<QUICStream>)> onStream) {
        QUICServerCredentials creds;
        creds.selfSigned = true;
        listener = std::make_unique<ListenerQUIC>("127.0.0.1", 0, alpn, creds, [this, onStream](ClientQUIC* peer) {
            peer->OnStream([this, onStream](QUICStream stream) {
                onStream(*this, std::make_shared<QUICStream>(std::move(stream)));
            });
        });
        listener->ListenAsyncAsync();
    }
    void Keep(std::shared_ptr<QUICStream> stream) {
        std::lock_guard lock(mutex);
        kept.push_back(std::move(stream));
    }
    std::uint16_t Port() const { return listener->Port(); }
};

} // namespace

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

    try {
        // ── What arrived first is delivered by OnRecieve ─────────────────────
        {
            Seen& seen = *new Seen;
            std::atomic<bool> registering = false;
            std::atomic<int> deliveredInside = 0;
            Server& server = *new Server([&](Server& self, std::shared_ptr<QUICStream> stream) {
                // Long enough for "hello" and its FIN to be waiting.
                std::this_thread::sleep_for(std::chrono::milliseconds(300));
                registering = true;
                auto onData = seen.OnData();
                auto onEnd = seen.OnEnd();
                stream->OnRecieve(
                    [&, onData](std::span<const char> data) { if (registering) ++deliveredInside; onData(data); },
                    [&, onEnd](bool clean) { if (registering) ++deliveredInside; onEnd(clean); });
                registering = false;
                self.Keep(stream);
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            stream.SendSync("hello", 5, true);
            Check(WaitFor([&] { return seen.Ends() == 1; }, std::chrono::seconds(10)) && seen.Bytes() == "hello" && seen.clean,
                  std::format("bytes and a FIN that came before OnRecieve are delivered by it: '{}', clean", seen.Bytes()));
            Check(deliveredInside.load() == 2, "on the registering thread, before OnRecieve returned");
        }

        // ── Registered mid-flow: one callback at a time, every byte in order ──
        {
            Seen& seen = *new Seen;
            auto slow = std::make_shared<std::atomic<int>>(50);
            Server& server = *new Server([&seen, slow](Server& self, std::shared_ptr<QUICStream> stream) {
                // Bytes keep coming meanwhile, so the backlog this thread
                // delivers and the worker's new arrivals overlap in time.
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
                auto onData = seen.OnData();
                std::weak_ptr<QUICStream> weak = stream;
                stream->OnRecieve(
                    [&seen, slow, onData, weak](std::span<const char> data) {
                        if (slow->fetch_sub(1) > 0) std::this_thread::sleep_for(std::chrono::microseconds(200));
                        onData(data);
                        // Echoed from inside the callback; its onSent is one
                        // of the stream's callbacks too.
                        if (auto s = weak.lock()) {
                            s->SendAsync(data.data(), static_cast<std::uint32_t>(data.size()), false, [&seen](bool) {
                                seen.Enter();
                                seen.Leave();
                            });
                        }
                    },
                    [weak, onEnd = seen.OnEnd()](bool clean) {
                        onEnd(clean);
                        if (auto s = weak.lock()) s->SendAsync(nullptr, 0, true, nullptr);
                    });
                self.Keep(stream);
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            Seen& echo = *new Seen;
            stream.OnRecieve(echo.OnData(), echo.OnEnd());
            constexpr std::size_t total = 2000 * 512;
            std::vector<char> chunk(512);
            for (std::size_t sent = 0; sent < total; sent += chunk.size()) {
                for (std::size_t i = 0; i < chunk.size(); ++i) chunk[i] = Pattern(sent + i);
                stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(chunk.size()), sent + chunk.size() == total, nullptr);
            }
            const bool ended = WaitFor([&] { return seen.Ends() == 1 && echo.Ends() == 1; }, std::chrono::seconds(20));
            Check(ended && seen.clean && FollowsPattern(seen.Bytes(), total) && seen.dataAfterEnd == 0,
                  std::format("registered mid-flow, every byte arrived in order, then a clean end ({} of {})", seen.Bytes().size(), total));
            Check(seen.overlaps.load() == 0, "and no two of the stream's callbacks, onSent included, ever ran at once");
            Check(echo.clean && FollowsPattern(echo.Bytes(), total), "the echo sent from inside onData came back whole");
        }

        // ── No pool ──────────────────────────────────────────────────────────
        {
            Seen& seen = *new Seen;
            std::atomic<bool> registered = false;
            Server& server = *new Server([&](Server& self, std::shared_ptr<QUICStream> stream) {
                std::weak_ptr<QUICStream> weak = stream;
                stream->OnRecieve(
                    [&, weak](std::span<const char> data) {
                        seen.OnData()(data);
                        if (auto s = weak.lock()) s->SendAsync(data.data(), static_cast<std::uint32_t>(data.size()), false, nullptr);
                    },
                    [weak](bool) { if (auto s = weak.lock()) s->SendAsync(nullptr, 0, true, nullptr); });
                self.Keep(stream);
                registered = true;
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            Seen& echo = *new Seen;
            stream.OnRecieve(echo.OnData(), echo.OnEnd());
            const char first = Pattern(0);
            stream.SendAsync(&first, 1, false, nullptr);
            // The server's OnStream is a pool task: let it run first.
            WaitFor([&] { return registered.load(); }, std::chrono::seconds(10));

            Parking& parking = *new Parking;
            const bool allParked = parking.Park();
            constexpr std::size_t total = 256 * 1024;
            std::vector<char> rest(total - 1);
            for (std::size_t i = 0; i < rest.size(); ++i) rest[i] = Pattern(i + 1);
            stream.SendAsync(rest.data(), static_cast<std::uint32_t>(rest.size()), true, nullptr);
            const bool ended = WaitFor([&] { return echo.Ends() == 1; }, std::chrono::seconds(10));
            Check(allParked && ended && echo.clean && FollowsPattern(echo.Bytes(), total),
                  std::format("with all {} pool workers parked, 256 KiB went out and came back by callback ({} bytes)",
                              parking.parked.load(), echo.Bytes().size()));
            parking.Release();
        }

        // ── Unclean ends ─────────────────────────────────────────────────────
        {
            std::mutex mutex;
            std::vector<Seen*> seens;
            Server& server = *new Server([&](Server& self, std::shared_ptr<QUICStream> stream) {
                Seen& seen = *new Seen;
                {
                    std::lock_guard lock(mutex);
                    seens.push_back(&seen);
                }
                stream->OnRecieve(seen.OnData(), seen.OnEnd());
                self.Keep(stream);
            });
            auto latest = [&]() -> Seen* {
                std::lock_guard lock(mutex);
                return seens.empty() ? nullptr : seens.back();
            };
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            {
                QUICStream stream = client.OpenStream();
                stream.SendSync("x", 1, false);
                WaitFor([&] { return latest() != nullptr && latest()->Bytes() == "x"; }, std::chrono::seconds(10));
                stream.Abort(7);
                Seen* seen = latest();
                Check(seen && WaitFor([&] { return seen->Ends() == 1; }, std::chrono::seconds(10)) && !seen->clean,
                      "the peer's Abort() ends the stream, not clean");
            }
            {
                ClientQUIC lost("127.0.0.1", server.Port(), alpn, ClientCredentials());
                QUICStream stream = lost.OpenStream();
                Seen* before = latest();
                stream.SendSync("y", 1, false);
                WaitFor([&] { return latest() != before && latest()->Bytes() == "y"; }, std::chrono::seconds(10));
                Seen* seen = latest();
                lost.Stop();   // the connection goes; the stream gets no FIN
                Check(seen != before && WaitFor([&] { return seen->Ends() == 1; }, std::chrono::seconds(10)) && !seen->clean,
                      "a lost connection ends it, not clean");
                std::this_thread::sleep_for(std::chrono::milliseconds(100));
                Check(seen->Ends() == 1, "each end exactly once");
            }
        }

        // ── A byte cap, by Abort() from inside onData ────────────────────────
        {
            Seen& seen = *new Seen;
            std::atomic<int> afterAbort = 0;
            std::atomic<bool> aborted = false;
            constexpr std::size_t cap = 64 * 1024;
            Server& server = *new Server([&](Server& self, std::shared_ptr<QUICStream> stream) {
                std::weak_ptr<QUICStream> weak = stream;
                stream->OnRecieve(
                    [&, weak](std::span<const char> data) {
                        if (aborted) ++afterAbort;
                        seen.OnData()(data);
                        if (seen.Bytes().size() > cap && !aborted) {
                            aborted = true;
                            if (auto s = weak.lock()) s->Abort(0x42);
                        }
                    },
                    seen.OnEnd());
                self.Keep(stream);
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            Seen& back = *new Seen;
            stream.OnRecieve(back.OnData(), back.OnEnd());
            constexpr int sends = 16;
            std::atomic<int> reported = 0;
            std::atomic<int> refused = 0;
            const std::vector<char> chunk(64 * 1024, 'z');
            for (int i = 0; i < sends; ++i) {
                stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(chunk.size()), false, [&](bool sent) {
                    if (!sent) ++refused;
                    ++reported;
                });
            }
            const bool ended = WaitFor([&] { return seen.Ends() == 1; }, std::chrono::seconds(10));
            Check(ended && !seen.clean && afterAbort.load() == 0 && seen.Bytes().size() > cap
                      && seen.Bytes().size() < sends * chunk.size(),
                  std::format("Abort() inside onData stopped the stream at {} KiB: no onData after it, onEnd(false)",
                              seen.Bytes().size() / 1024));
            // Its sends may all report true: msquic takes them into its send
            // buffer before the reset arrives, and true means just that. A
            // send made once the reset has arrived is the one refused.
            const bool backEnded = WaitFor([&] { return back.Ends() == 1 && reported.load() == sends; }, std::chrono::seconds(10));
            bool probeRefused = false;
            for (int attempt = 0; attempt < 100 && !probeRefused; ++attempt) {
                auto result = std::make_shared<std::atomic<int>>(-1);
                stream.SendAsync("p", 1, false, [result](bool sent) { *result = sent ? 1 : 0; });
                WaitFor([&] { return result->load() != -1; }, std::chrono::seconds(5));
                probeRefused = result->load() == 0;
                if (!probeRefused) std::this_thread::sleep_for(std::chrono::milliseconds(20));
            }
            Check(backEnded && !back.clean && probeRefused,
                  std::format("the peer saw both directions reset: its receive side ended not clean, and its sends are refused ({} of the first {} already were)",
                              refused.load(), sends));
        }

        // ── The wrapper destroyed or moved under its callbacks ───────────────
        {
            // Destroyed inside its own onEnd: the callback holds the last
            // reference.
            std::atomic<int> destroyedInside = 0;
            std::atomic<int> streams = 0;
            Server& server = *new Server([&](Server&, std::shared_ptr<QUICStream> stream) {
                ++streams;
                auto holder = std::make_shared<std::shared_ptr<QUICStream>>(stream);
                stream->OnRecieve([](std::span<const char>) {}, [holder, &destroyedInside](bool) {
                    holder->reset();
                    ++destroyedInside;
                });
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            for (int i = 0; i < 3; ++i) {
                QUICStream stream = client.OpenStream();
                stream.SendSync("bye", 3, true);
            }
            Check(WaitFor([&] { return destroyedInside.load() == 3; }, std::chrono::seconds(10)),
                  "a wrapper destroyed inside its own onEnd breaks nothing, three times over");
        }
        {
            // Moved while bytes flow: the callbacks never touch the wrapper.
            Seen& seen = *new Seen;
            std::atomic<bool> registered = false;
            std::shared_ptr<QUICStream> first;
            Server& server = *new Server([&](Server&, std::shared_ptr<QUICStream> stream) {
                stream->OnRecieve(seen.OnData(), seen.OnEnd());
                first = stream;
                registered = true;
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            const char lead = Pattern(0);
            stream.SendSync(&lead, 1, false);
            WaitFor([&] { return registered.load(); }, std::chrono::seconds(10));
            constexpr std::size_t total = 512 * 1024;
            std::thread sender([&] {
                std::vector<char> chunk(1024);
                for (std::size_t sent = 1; sent < total; sent += chunk.size()) {
                    const std::size_t size = std::min(chunk.size(), total - sent);
                    for (std::size_t i = 0; i < size; ++i) chunk[i] = Pattern(sent + i);
                    stream.SendSync(chunk.data(), static_cast<std::uint32_t>(size), sent + size == total);
                }
            });
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
            auto second = std::make_shared<QUICStream>(std::move(*first));
            first.reset();
            sender.join();
            Check(WaitFor([&] { return seen.Ends() == 1; }, std::chrono::seconds(10)) && seen.clean
                      && FollowsPattern(seen.Bytes(), total),
                  "a wrapper moved while bytes flow keeps receiving, every byte in order");
            second.reset();
        }
        {
            // Destroyed from another thread mid-flow: no callback after.
            // What the callbacks touch is leaked, in case they do run late.
            struct Watch {
                Seen seen;
                std::atomic<bool> destroyed = false;
                std::atomic<int> afterDestroyed = 0;
                std::shared_ptr<QUICStream> held;
                std::atomic<bool> registered = false;
            };
            Watch& watch = *new Watch;
            Seen& seen = watch.seen;
            std::atomic<bool>& destroyed = watch.destroyed;
            std::atomic<int>& afterDestroyed = watch.afterDestroyed;
            std::shared_ptr<QUICStream>& held = watch.held;
            std::atomic<bool>& registered = watch.registered;
            Server& server = *new Server([&watch](Server&, std::shared_ptr<QUICStream> stream) {
                stream->OnRecieve(
                    [&watch](std::span<const char>) {
                        watch.seen.Enter();
                        // Slow, so destroying the wrapper often catches one.
                        std::this_thread::sleep_for(std::chrono::microseconds(500));
                        if (watch.destroyed) ++watch.afterDestroyed;
                        watch.seen.Leave();
                    },
                    [&watch](bool) { if (watch.destroyed) ++watch.afterDestroyed; });
                watch.held = stream;
                watch.registered = true;
            });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            stream.SendSync("go", 2, false);
            WaitFor([&] { return registered.load(); }, std::chrono::seconds(10));
            std::atomic<bool> stop = false;
            std::thread sender([&] {
                const std::vector<char> chunk(256, 'q');
                while (!stop) {
                    stream.SendAsync(chunk.data(), static_cast<std::uint32_t>(chunk.size()), false, nullptr);
                    std::this_thread::sleep_for(std::chrono::microseconds(200));
                }
            });
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
            held.reset();       // ~QUICStream on this thread, callbacks running on the worker
            destroyed = true;
            std::this_thread::sleep_for(std::chrono::milliseconds(200));
            stop = true;
            sender.join();
            Check(afterDestroyed.load() == 0 && seen.overlaps.load() == 0,
                  "once a wrapper destroyed on another thread has gone, no callback of it runs again");
        }

        // ── Misuse and edge cases ────────────────────────────────────────────
        {
            Server& server = *new Server([](Server& self, std::shared_ptr<QUICStream> stream) { self.Keep(stream); });
            ClientQUIC client("127.0.0.1", server.Port(), alpn, ClientCredentials());
            QUICStream stream = client.OpenStream();
            stream.OnRecieve([](std::span<const char>) {}, [](bool) {});
            bool threw = false;
            try {
                stream.OnRecieve([](std::span<const char>) {}, [](bool) {});
            } catch (const QUICException&) {
                threw = true;
            }
            Check(threw, "a second OnRecieve on one stream throws QUICException");

            QUICStream sendOnly = client.OpenStream(/*unidirectional=*/true);
            std::atomic<int> sendOnlyEnds = 0;
            std::atomic<bool> sendOnlyClean = true;
            sendOnly.OnRecieve([](std::span<const char>) {}, [&](bool clean) { ++sendOnlyEnds; sendOnlyClean = clean; });
            Check(sendOnlyEnds.load() == 1 && !sendOnlyClean.load(), "a stream that cannot receive gets onEnd(false) at once");

            QUICStream none;
            std::atomic<int> noneEnds = 0;
            none.OnRecieve([](std::span<const char>) {}, [&](bool clean) { if (!clean) ++noneEnds; });
            Check(noneEnds.load() == 1, "so does no stream at all");
        }
    } 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 peers' idle timeout. _Exit doesn't flush.
    std::fflush(stdout);
    std::_Exit(failures == 0 ? 0 : 1);
}