// quic_incremental_dispatch_test.cpp
//
// Regression test for setIncrementalStreamDispatch(true) + its recv_buffer
// trimming (recv_base_offset): sends a large payload on one bidi stream to a
// server that dispatches each newly-contiguous chunk as it arrives instead
// of buffering the whole message, accumulates the chunks server-side,
// verifies the reassembled data's SHA-256 matches the original, then echoes
// it back and verifies the client's received bytes match too. Also asserts
// more than one non-fin chunk was actually delivered, so the trimming path
// (recv_base_offset advancing mid-stream) is genuinely exercised more than
// once per stream, not just at FIN.

#include <iostream>
#include <string>
#include <vector>
#include <map>
#include <mutex>
#include <thread>
#include <atomic>
#include <chrono>
#include <stdexcept>

#include "connection.h"
#include "eventapi.h"
#include "socket.h"
#include "exception.h"
#include "random.h"
#include "crypto/sha.h"

#include "https_certs.h"
#include "https_ca_cert.h"

using namespace netplus;

static int g_passed = 0, g_failed = 0;

static void check(bool ok, const std::string& name) {
    if (ok) { std::cout << "  PASS: " << name << std::endl; g_passed++; }
    else    { std::cout << "  FAIL: " << name << std::endl; g_failed++; }
}

static std::string toHex(const std::vector<uint8_t>& v) {
    static const char* hexd = "0123456789abcdef";
    std::string s;
    s.reserve(v.size() * 2);
    for (uint8_t b : v) {
        s.push_back(hexd[b >> 4]);
        s.push_back(hexd[b & 0x0F]);
    }
    return s;
}

// ============================================================================
// QUIC server: incremental-dispatch echo. Accumulates chunks per stream_id
// as they're pushed (rather than getting the whole message once at FIN),
// and echoes the full reassembled message back once FIN arrives.
// ============================================================================

static std::mutex g_accum_mutex;
static std::map<uint64_t, std::vector<uint8_t>> g_accum;
static std::atomic<int> g_nonfin_chunk_count{0};
static std::atomic<int> g_fin_chunk_count{0};

class IncrementalEchoServer : public event {
public:
    IncrementalEchoServer(std::vector<netplus::socket*> socks, int timeout = 500)
        : event(socks, timeout) {}

    void RequestEvent(con&, const int, ULONG_PTR) override {}
    void ResponseEvent(con&, const int, ULONG_PTR) override {}
    void ConnectEvent(con&, const int, ULONG_PTR) override {}
    void DisconnectEvent(con&, const int, ULONG_PTR) override {}
    void CreateConnection(std::shared_ptr<con>& res) override {
        res = std::make_shared<con>(this);
    }
};

static std::atomic<bool> g_quic_ready(false);

static void waitReady(std::atomic<bool>& ready) {
    while (!ready.load())
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
}

static void runQuicServer(std::map<std::string, ssl::CertificateBundle>& certs, int port) {
    try {
        quic serverSock(certs, "127.0.0.1", port, 64, -1);
        serverSock.setIncrementalStreamDispatch(true);
        serverSock.setStreamCallback([](netplus::socket* sock, uint64_t stream_id,
                                        const std::vector<uint8_t>& data, bool fin) {
            netplus::quic* q = dynamic_cast<netplus::quic*>(sock);
            if (!q) return;

            if (fin) g_fin_chunk_count++;
            else if (!data.empty()) g_nonfin_chunk_count++;

            std::vector<uint8_t> whole;
            {
                std::lock_guard<std::mutex> lk(g_accum_mutex);
                auto& buf = g_accum[stream_id];
                buf.insert(buf.end(), data.begin(), data.end());
                if (fin) whole = buf;
            }
            if (fin) {
                q->sendStreamData(stream_id, whole, true);
            }
        });

        IncrementalEchoServer srv({&serverSock});
        g_quic_ready.store(true);
        srv.runEventloop();
    } catch (NetException& e) {
        std::cerr << "[QUIC server] NetException: " << e.what() << std::endl;
        g_quic_ready.store(true);
    } catch (std::exception& e) {
        std::cerr << "[QUIC server] Exception: " << e.what() << std::endl;
        g_quic_ready.store(true);
    }
}

static std::vector<uint8_t> quicRoundTrip(const std::vector<uint8_t>& payload,
                                           const std::string& host, int port) {
    quic client;
    client.connect(host, port);
    uint64_t sid = client.openStream(true);

    client.sendStreamData(sid, payload, true);

    std::vector<uint8_t> out(payload.size());
    size_t total = 0;
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(60);
    while (total < out.size() && std::chrono::steady_clock::now() < deadline) {
        client.pumpNetwork(MSG_DONTWAIT);
        if (client.hasStreamData(sid)) {
            size_t got = client.recvStreamData(sid, out.data() + total, out.size() - total);
            if (got > 0) { total += got; continue; }
        }
        std::this_thread::sleep_for(std::chrono::microseconds(200));
    }
    if (total != out.size())
        throw std::runtime_error("quicRoundTrip: incomplete echo (" +
                                  std::to_string(total) + "/" + std::to_string(out.size()) + " bytes)");
    return out;
}

int main() {
    std::cout << "=== QUIC Incremental Stream Dispatch Test ===" << std::endl;

    const size_t kPayloadSize = 4096 * 1024; // 4 MB: forces many dispatch/trim cycles
    const int kQuicPort = 19549;

    int rc = 0;
    try {
        std::vector<uint8_t> original(kPayloadSize);
        fillRandomBytes(original.data(), original.size());
        std::vector<uint8_t> originalHash = sha256_hash(original);
        std::cout << "Generated " << kPayloadSize << " random bytes, sha256=" << toHex(originalHash) << std::endl;

        x509cert cert;
        if (!cert.loadFromBuffer(test_cert_der)) {
            std::cerr << "Failed to load certificate" << std::endl;
            return 1;
        }
        std::map<std::string, ssl::CertificateBundle> certs;
        ssl::CertificateBundle bundle;
        bundle.cert = cert;
        bundle.privateKeyDer = std::vector<uint8_t>(test_key_der.begin(), test_key_der.end());
        bundle.rsa_key = rsa(bundle.privateKeyDer);
        bundle.chain.push_back(std::vector<uint8_t>(MKCERT_ROOT_CA_DER,
            MKCERT_ROOT_CA_DER + MKCERT_ROOT_CA_DER_LEN));
        certs["localhost"] = bundle;
        certs["127.0.0.1"] = bundle;

        std::thread quicServerThread(runQuicServer, std::ref(certs), kQuicPort);
        waitReady(g_quic_ready);

        std::vector<uint8_t> echoed = quicRoundTrip(original, "127.0.0.1", kQuicPort);

        event::Running = false;
        quicServerThread.join();

        check(echoed == original, "incremental-dispatch echo preserves bytes exactly");
        check(sha256_hash(echoed) == originalHash, "incremental-dispatch echo preserves sha256");
        check(g_nonfin_chunk_count.load() > 1,
              "server received more than one non-fin chunk (trim path exercised repeatedly, not just at FIN): got " +
              std::to_string(g_nonfin_chunk_count.load()));
        check(g_fin_chunk_count.load() == 1, "server received exactly one fin=true chunk");

        std::cout << "non-fin chunks: " << g_nonfin_chunk_count.load()
                  << ", fin chunks: " << g_fin_chunk_count.load() << std::endl;

    } catch (NetException& e) {
        std::cerr << "NetException: " << e.what() << std::endl;
        rc = 1;
    } catch (std::exception& e) {
        std::cerr << "Exception: " << e.what() << std::endl;
        rc = 1;
    }

    std::cout << "\n==============================" << std::endl;
    std::cout << "Results: " << g_passed << " passed, " << g_failed << " failed" << std::endl;
    std::cout << "==============================" << std::endl;

    return (rc != 0 || g_failed > 0) ? 1 : 0;
}
