/*******************************************************************************
Copyright (c) 2026, Jan Koester jan.koester@gmx.net
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the <organization> nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*******************************************************************************/

// HTTP/2 checksum integrity test, exercised two ways against the SAME
// server (one HttpEvent instance listening on both a plain-TCP and a
// TLS+ALPN socket):
//
//  1. Hand-rolled real HTTP/2 wire frames (HPACK-encoded HEADERS + DATA,
//     RFC 7540 flow control via WINDOW_UPDATE) over cleartext h2c, the same
//     style as httpd_h2_concurrency_test.cpp -- this is what lets the test
//     also exercise the library's true streaming-ingestion callback API
//     (onH2StreamHeaders/onH2DataChunk/sendH2StreamResponse) directly,
//     which HttpClient's own H2 client never touches.
//  2. libhttppp's own HttpClient (vers=2) over real TLS with ALPN
//     negotiating "h2", i.e. the library's actual production H2 client
//     path (_h2Request()) -- HttpClient never speaks cleartext h2c, so
//     this needs the TLS listener.
//
// Both drive the same size sweep (empty to 1MB) for GET (server-generated,
// checksummed via X-Sha256) and POST (echoed back, checksummed both ways).

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <thread>
#include <utility>

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#include <netplus/exception.h>
#include <netplus/socket.h>
#include <netplus/crypto/tls.h>
#include <netplus/crypto/x509.h>

#include "checksum_common.h"
#include "hpack.h"
#include "https_certs.h"
#include "https_ca_cert.h"

using namespace checksumtest;

namespace {

int g_passed = 0, g_failed = 0;

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

constexpr uint8_t FRAME_DATA          = 0x00;
constexpr uint8_t FRAME_HEADERS       = 0x01;
constexpr uint8_t FRAME_SETTINGS      = 0x04;
constexpr uint8_t FRAME_PING          = 0x06;
constexpr uint8_t FRAME_GOAWAY        = 0x07;
constexpr uint8_t FRAME_WINDOW_UPDATE = 0x08;

constexpr uint8_t FLAG_ACK         = 0x01;
constexpr uint8_t FLAG_END_STREAM  = 0x01;
constexpr uint8_t FLAG_END_HEADERS = 0x04;

constexpr size_t FRAME_HEADER_LEN = 9;
constexpr size_t MAX_FRAME_SIZE = 16384;
constexpr char CLIENT_PREFACE[] = "PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n";

// ---------------------------------------------------------------------
// Server: adds a true-streaming "/stream" POST path (exercises
// onH2StreamHeaders/onH2DataChunk/sendH2StreamResponse) on top of the
// protocol-agnostic ChecksumEvent handler shared with the H1/H3 tests.
// ---------------------------------------------------------------------
class H2ChecksumEvent : public ChecksumEvent {
public:
    using ChecksumEvent::ChecksumEvent;

    bool onH2StreamHeaders(libhttppp::HttpRequest &, uint32_t streamId,
                            const std::vector<libhttppp::hpack::HeaderField> &headers) override {
        std::string path;
        for (auto &h : headers) {
            if (h.name == ":path") { path = h.value; break; }
        }
        if (path != "/stream") return false;

        std::lock_guard<std::mutex> lk(_mu);
        _streamBodies[streamId].clear();
        return true;
    }

    void onH2DataChunk(libhttppp::HttpRequest &, uint32_t streamId,
                        const char *data, size_t len, bool endStream,
                        std::string &h2out, const int, ULONG_PTR) override {
        std::string finishedBody;
        bool finished = false;
        {
            std::lock_guard<std::mutex> lk(_mu);
            auto it = _streamBodies.find(streamId);
            if (it == _streamBodies.end()) return;
            if (data && len) it->second.append(data, len);
            if (endStream) {
                finishedBody = std::move(it->second);
                _streamBodies.erase(it);
                finished = true;
            }
        }
        if (finished) {
            std::string hash = sha256Hex(finishedBody.data(), finishedBody.size());
            sendH2StreamResponse(h2out, streamId, 200, "text/plain", hash);
        }
    }

private:
    std::mutex _mu;
    std::map<uint32_t, std::string> _streamBodies;
};

netplus::ssl::CertificateBundle buildTestCertBundle() {
    netplus::x509cert cert;
    if (!cert.loadFromBuffer(test_cert_der))
        throw std::runtime_error("failed to load embedded test certificate");

    netplus::ssl::CertificateBundle bundle;
    bundle.cert = cert;
    bundle.privateKeyDer = std::vector<uint8_t>(test_key_der.begin(), test_key_der.end());
    bundle.rsa_key = netplus::rsa(bundle.privateKeyDer);
    bundle.chain.push_back(std::vector<uint8_t>(
        MKCERT_ROOT_CA_DER, MKCERT_ROOT_CA_DER + MKCERT_ROOT_CA_DER_LEN));
    return bundle;
}

// Starts one H2ChecksumEvent listening on BOTH a plain-TCP socket (for the
// hand-rolled cleartext h2c client below) and a TLS+ALPN socket (for
// HttpClient's real h2 client, which only ever negotiates h2 over TLS).
// Returns {h2c port, TLS port}.
std::pair<int, int> startServer(std::unique_ptr<netplus::tcp> &sockOwner,
                                 std::unique_ptr<netplus::ssl> &sslSockOwner,
                                 std::unique_ptr<H2ChecksumEvent> &serverOwner,
                                 std::thread &serverThread,
                                 int tlsPort) {
    sockOwner = std::make_unique<netplus::tcp>("127.0.0.1", 0, 16, -1);
    sockOwner->bind();

    sockaddr_in bound{};
    socklen_t blen = sizeof(bound);
    ::getsockname(sockOwner->fd(), reinterpret_cast<sockaddr *>(&bound), &blen);
    int port = ntohs(bound.sin_port);

    sockOwner->listen();

    std::map<std::string, netplus::ssl::CertificateBundle> certs;
    netplus::ssl::CertificateBundle bundle = buildTestCertBundle();
    certs["localhost"] = bundle;
    certs["127.0.0.1"] = bundle;
    sslSockOwner = std::make_unique<netplus::ssl>(certs, "127.0.0.1", tlsPort, 16, -1);

    serverOwner = std::make_unique<H2ChecksumEvent>(
        std::vector<netplus::socket *>{sockOwner.get(), sslSockOwner.get()});
    serverThread = std::thread([s = serverOwner.get()] { s->runEventloop(); });
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
    return {port, tlsPort};
}

// ---------------------------------------------------------------------
// Minimal real HTTP/2 client: hand-built frames, real HPACK encode/decode,
// real RFC 7540 flow control (replenishes WINDOW_UPDATE for every DATA
// frame received, both stream- and connection-level, exactly like a
// conformant client must to receive bodies past the default 64KB window).
// ---------------------------------------------------------------------
class H2RawClient {
public:
    struct Response {
        int status = 0;
        std::map<std::string, std::string> headers;
        std::vector<char> body;
    };

    void connectTo(int port) {
        _fd = ::socket(AF_INET, SOCK_STREAM, 0);
        sockaddr_in addr{};
        addr.sin_family = AF_INET;
        addr.sin_port = htons(static_cast<uint16_t>(port));
        inet_pton(AF_INET, "127.0.0.1", &addr.sin_addr);
        if (::connect(_fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) < 0) {
            ::close(_fd);
            throw std::runtime_error("H2RawClient: connect failed");
        }
        sendAll(std::string(CLIENT_PREFACE, sizeof(CLIENT_PREFACE) - 1));
        sendFrame(FRAME_SETTINGS, 0, 0, ""); // empty client SETTINGS
    }

    ~H2RawClient() { if (_fd >= 0) ::close(_fd); }

    Response request(const std::string &method, const std::string &path,
                      const std::vector<char> &body) {
        uint32_t sid = _nextStreamId;
        _nextStreamId += 2;

        std::string hpackOut = libhttppp::hpack::Encoder::encodeRequestHeaders(
            method, path, "http", "127.0.0.1");

        bool hasBody = !body.empty();
        sendFrame(FRAME_HEADERS,
                   FLAG_END_HEADERS | (hasBody ? 0 : FLAG_END_STREAM),
                   sid, hpackOut);

        size_t off = 0;
        while (off < body.size()) {
            size_t take = std::min(MAX_FRAME_SIZE, body.size() - off);
            bool last = (off + take == body.size());
            sendFrame(FRAME_DATA, last ? FLAG_END_STREAM : 0, sid,
                       std::string(body.data() + off, take));
            off += take;
        }

        Response res;
        for (;;) {
            Frame f = readFrame();

            if (f.streamId == 0) {
                if (f.type == FRAME_SETTINGS && !(f.flags & FLAG_ACK)) {
                    sendFrame(FRAME_SETTINGS, FLAG_ACK, 0, "");
                } else if (f.type == FRAME_PING && !(f.flags & FLAG_ACK) &&
                           f.payload.size() == 8) {
                    sendFrame(FRAME_PING, FLAG_ACK, 0, f.payload);
                } else if (f.type == FRAME_GOAWAY) {
                    throw std::runtime_error("H2RawClient: server sent GOAWAY");
                }
                continue;
            }

            if (f.streamId != sid) continue; // unrelated stream traffic

            if (f.type == FRAME_HEADERS) {
                auto decoded = _decoder.decode(
                    reinterpret_cast<const uint8_t *>(f.payload.data()), f.payload.size());
                for (auto &h : decoded) {
                    if (h.name == ":status") {
                        try { res.status = std::stoi(h.value); } catch (...) {}
                    } else {
                        res.headers[h.name] = h.value;
                    }
                }
                if (f.flags & FLAG_END_STREAM) break;
            } else if (f.type == FRAME_DATA) {
                if (!f.payload.empty()) {
                    res.body.insert(res.body.end(), f.payload.begin(), f.payload.end());
                    sendFrame(FRAME_WINDOW_UPDATE, 0, sid, windowIncrement(f.payload.size()));
                    sendFrame(FRAME_WINDOW_UPDATE, 0, 0, windowIncrement(f.payload.size()));
                }
                if (f.flags & FLAG_END_STREAM) break;
            }
        }
        return res;
    }

private:
    struct Frame {
        uint8_t type = 0;
        uint8_t flags = 0;
        uint32_t streamId = 0;
        std::string payload;
    };

    static std::string windowIncrement(size_t n) {
        uint32_t v = static_cast<uint32_t>(n) & 0x7FFFFFFFu;
        std::string p(4, '\0');
        p[0] = static_cast<char>((v >> 24) & 0xFF);
        p[1] = static_cast<char>((v >> 16) & 0xFF);
        p[2] = static_cast<char>((v >> 8) & 0xFF);
        p[3] = static_cast<char>(v & 0xFF);
        return p;
    }

    void sendFrame(uint8_t type, uint8_t flags, uint32_t streamId, const std::string &payload) {
        std::string out;
        out.reserve(FRAME_HEADER_LEN + payload.size());
        uint32_t len = static_cast<uint32_t>(payload.size());
        out.push_back(static_cast<char>((len >> 16) & 0xFF));
        out.push_back(static_cast<char>((len >> 8) & 0xFF));
        out.push_back(static_cast<char>(len & 0xFF));
        out.push_back(static_cast<char>(type));
        out.push_back(static_cast<char>(flags));
        out.push_back(static_cast<char>((streamId >> 24) & 0x7F));
        out.push_back(static_cast<char>((streamId >> 16) & 0xFF));
        out.push_back(static_cast<char>((streamId >> 8) & 0xFF));
        out.push_back(static_cast<char>(streamId & 0xFF));
        out += payload;
        sendAll(out);
    }

    void sendAll(const std::string &data) {
        size_t sent = 0;
        while (sent < data.size()) {
            ssize_t n = ::send(_fd, data.data() + sent, data.size() - sent, 0);
            if (n <= 0) throw std::runtime_error("H2RawClient: send failed");
            sent += static_cast<size_t>(n);
        }
    }

    void fillAtLeast(size_t n) {
        auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(15);
        while (_inbuf.size() < n) {
            if (std::chrono::steady_clock::now() >= deadline)
                throw std::runtime_error("H2RawClient: timeout waiting for frame data");
            char tmp[65536];
            ssize_t r = ::recv(_fd, tmp, sizeof(tmp), 0);
            if (r <= 0) throw std::runtime_error("H2RawClient: connection closed unexpectedly");
            _inbuf.insert(_inbuf.end(), tmp, tmp + r);
        }
    }

    Frame readFrame() {
        fillAtLeast(FRAME_HEADER_LEN);
        uint32_t len = (static_cast<uint8_t>(_inbuf[0]) << 16) |
                       (static_cast<uint8_t>(_inbuf[1]) << 8) |
                       static_cast<uint8_t>(_inbuf[2]);
        Frame f;
        f.type = static_cast<uint8_t>(_inbuf[3]);
        f.flags = static_cast<uint8_t>(_inbuf[4]);
        f.streamId = ((static_cast<uint8_t>(_inbuf[5]) & 0x7F) << 24) |
                     (static_cast<uint8_t>(_inbuf[6]) << 16) |
                     (static_cast<uint8_t>(_inbuf[7]) << 8) |
                     static_cast<uint8_t>(_inbuf[8]);
        fillAtLeast(FRAME_HEADER_LEN + len);
        f.payload.assign(_inbuf.begin() + FRAME_HEADER_LEN,
                          _inbuf.begin() + FRAME_HEADER_LEN + len);
        _inbuf.erase(_inbuf.begin(), _inbuf.begin() + FRAME_HEADER_LEN + len);
        return f;
    }

    int _fd = -1;
    uint32_t _nextStreamId = 1;
    libhttppp::hpack::Decoder _decoder;
    std::vector<char> _inbuf;
};

} // namespace

std::string headerValue(const libhttppp::HttpResponse &res, const std::string &key) {
    auto *hd = res.getHeaderData(key);
    if (hd && hd->getfirstValue())
        return hd->getfirstValue()->getvalue();
    return {};
}

int main() {
    std::cout << "=== HTTP/2 Checksum Integrity Test ===" << std::endl;
    int rc = 0;
    const int kTlsPort = 19644;

    try {
        std::unique_ptr<netplus::tcp> sock;
        std::unique_ptr<netplus::ssl> sslSock;
        std::unique_ptr<H2ChecksumEvent> server;
        std::thread serverThread;
        auto ports = startServer(sock, sslSock, server, serverThread, kTlsPort);
        ServerThreadGuard shutdownGuard{&serverThread};
        int port = ports.first;
        std::cout << "Server listening on 127.0.0.1:" << port
                  << " (h2c) and :" << kTlsPort << " (TLS+ALPN h2)" << std::endl;

        // ==================================================================
        // 1) Hand-rolled cleartext h2c client (real frames, real HPACK,
        //    real flow control) -- also exercises the streaming-ingestion
        //    callback API below, which HttpClient never touches.
        // ==================================================================
        H2RawClient clt;
        clt.connectTo(port);

        // ---- GET/POST size sweep, all on one connection (multiplexed
        // stream ids 1,3,5,... exercising real connection reuse) ----
        for (size_t size : sweepSizes()) {
            std::string label = "size=" + std::to_string(size);

            {
                auto res = clt.request("GET", "/gen/" + std::to_string(size), {});
                std::string got = sha256Hex(res.body);
                check(res.status == 200, "GET " + label + ": 200 response");
                check(res.body.size() == size, "GET " + label + ": body size matches");
                check(got == res.headers["x-sha256"],
                      "GET " + label + ": sha256 matches x-sha256 header");
            }

            {
                std::vector<char> payload = randomBuffer(size);
                std::string expectedHash = sha256Hex(payload);
                auto res = clt.request("POST", "/echo", payload);
                check(res.status == 200, "POST " + label + ": 200 response");
                check(res.body == payload, "POST " + label + ": echoed bytes match exactly");
                check(res.headers["x-sha256"] == expectedHash,
                      "POST " + label + ": server-reported sha256 matches client's");
            }
        }

        // ---- True streaming ingestion via onH2StreamHeaders/onH2DataChunk ----
        {
            std::vector<char> payload = randomBuffer(1048576);
            std::string expectedHash = sha256Hex(payload);
            auto res = clt.request("POST", "/stream", payload);
            std::string got(res.body.begin(), res.body.end());
            check(res.status == 200, "streaming ingestion: 200 response");
            check(got == expectedHash,
                  "streaming ingestion: server-hashed body (via onH2DataChunk) matches client's");
        }

        // ==================================================================
        // 2) libhttppp's own HttpClient (vers=2) over real TLS, negotiating
        //    "h2" via ALPN -- the library's actual production H2 client
        //    path (_h2Request()), against the SAME server.
        // ==================================================================
        {
            libhttppp::HttpUrl url("https://127.0.0.1:" + std::to_string(kTlsPort) + "/");
            libhttppp::HttpClient httpClt(url, /*vers=*/2, /*timeoutSec=*/5);

            for (size_t size : sweepSizes()) {
                std::string label = "size=" + std::to_string(size);

                {
                    libhttppp::HttpRequest req;
                    req.setRequestURL("/gen/" + std::to_string(size));
                    std::vector<char> body = httpClt.Get(req, /*maxTries=*/3);
                    std::string want = headerValue(*httpClt.lastResponse(), "x-sha256");
                    std::string got = sha256Hex(body);
                    check(httpClt.lastStatusCode() == 200,
                          "HttpClient GET " + label + ": 200 response");
                    check(body.size() == size,
                          "HttpClient GET " + label + ": body size matches");
                    check(got == want,
                          "HttpClient GET " + label + ": sha256 matches x-sha256 header");
                }

                {
                    std::vector<char> payload = randomBuffer(size);
                    std::string expectedHash = sha256Hex(payload);

                    libhttppp::HttpRequest req;
                    req.setRequestURL("/echo");
                    std::vector<char> body = httpClt.Post(req, payload, /*maxTries=*/3);
                    std::string reportedHash = headerValue(*httpClt.lastResponse(), "x-sha256");

                    check(httpClt.lastStatusCode() == 200,
                          "HttpClient POST " + label + ": 200 response");
                    check(body == payload,
                          "HttpClient POST " + label + ": echoed bytes match exactly");
                    check(reportedHash == expectedHash,
                          "HttpClient POST " + label + ": server-reported sha256 matches client's");
                }
            }

            // ---- Streaming download via GetStream()/readBodyChunk() and
            // readBodyChunkNonBlocking(), both over the real H2 client path
            // (_streamReadH2()) -- 1MB guarantees the response spans several
            // real DATA frames (H2C_MAX_FRAME_SIZE=16384) and at least one
            // WINDOW_UPDATE round trip, unlike the buffered Get() sweep above
            // which never touches this code at all. ----
            {
                const size_t size = 1048576;
                libhttppp::HttpRequest req;
                req.setRequestURL("/gen/" + std::to_string(size));
                libhttppp::HttpResponse res = httpClt.GetStream(req);
                std::string want = headerValue(res, "x-sha256");

                std::vector<char> collected;
                collected.reserve(size);
                char buf[4096];
                for (;;) {
                    size_t n = httpClt.readBodyChunk(buf, sizeof(buf));
                    if (n == 0) break;
                    collected.insert(collected.end(), buf, buf + n);
                }
                std::string got = sha256Hex(collected);
                check(res.getStatusCode() == 200, "HTTP/2 streaming GET (blocking): 200 response");
                check(collected.size() == size, "HTTP/2 streaming GET (blocking): total bytes match");
                check(got == want, "HTTP/2 streaming GET (blocking): incrementally-hashed sha256 matches header");
            }

            {
                const size_t size = 1048576;
                libhttppp::HttpRequest req;
                req.setRequestURL("/gen/" + std::to_string(size));
                libhttppp::HttpResponse res = httpClt.GetStream(req);
                std::string want = headerValue(res, "x-sha256");

                std::vector<char> collected;
                collected.reserve(size);
                char buf[4096];
                auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10);
                for (;;) {
                    if (std::chrono::steady_clock::now() > deadline) break;
                    size_t n = httpClt.readBodyChunkNonBlocking(buf, sizeof(buf));
                    if (n == static_cast<size_t>(-1)) break;
                    if (n == 0) {
                        std::this_thread::sleep_for(std::chrono::milliseconds(1));
                        continue;
                    }
                    collected.insert(collected.end(), buf, buf + n);
                }
                std::string got = sha256Hex(collected);
                check(res.getStatusCode() == 200, "HTTP/2 streaming GET (non-blocking): 200 response");
                check(collected.size() == size, "HTTP/2 streaming GET (non-blocking): total bytes match");
                check(got == want, "HTTP/2 streaming GET (non-blocking): incrementally-hashed sha256 matches header");
            }
        }

        // shutdownGuard stops and joins the server thread here (its
        // destructor runs before sock/sslSock/server as it's declared after
        // them).

    } catch (libhttppp::HTTPException &e) {
        std::cerr << "HTTPException: " << e.what() << std::endl;
        rc = 1;
    } catch (netplus::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 << "==============================" << 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;
}
