// quic_version_negotiation_test.cpp
//
// Exercises the RFC 9000 §6 / §17.2.1 Version Negotiation path: a
// long-header packet carrying a version the server doesn't speak must get
// a well-formed Version Negotiation packet back, not silence. Crafted with
// a raw UDP socket (not netplus::quic) so this only depends on the public
// wire format, not any internal/private server API.

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

#ifdef _WIN32
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <winsock2.h>
#include <ws2tcpip.h>
#ifdef _MSC_VER
#pragma comment(lib, "Ws2_32.lib")
#endif
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <poll.h>
#endif

#include "connection.h"
#include "eventapi.h"
#include "socket.h"
#include "exception.h"
#include "random.h"

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

using namespace netplus;

#ifdef _WIN32
using native_socket_t = SOCKET;
static constexpr native_socket_t kInvalidSocket = INVALID_SOCKET;
#else
using native_socket_t = int;
static constexpr native_socket_t kInvalidSocket = -1;
#endif

static void closeNativeSocket(native_socket_t fd) noexcept {
#ifdef _WIN32
    if (fd != kInvalidSocket) ::closesocket(fd);
#else
    if (fd != kInvalidSocket) ::close(fd);
#endif
}

static bool waitForReadable(native_socket_t fd, int timeoutMs) {
#ifdef _WIN32
    WSAPOLLFD pfd{};
    pfd.fd = fd;
    pfd.events = POLLRDNORM;
    const int rc = ::WSAPoll(&pfd, 1, timeoutMs);
    return rc > 0 && (pfd.revents & POLLRDNORM) != 0;
#else
    pollfd pfd{fd, POLLIN, 0};
    const int rc = ::poll(&pfd, 1, timeoutMs);
    return rc > 0 && (pfd.revents & POLLIN) != 0;
#endif
}

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

class NullEvents : public event {
public:
    NullEvents(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);
        NullEvents 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);
    }
}

int main() {
#ifdef _WIN32
    WSADATA wsaData{};
    const int wsaRc = ::WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (wsaRc != 0) {
        std::cerr << "WSAStartup failed: " << wsaRc << std::endl;
        return 1;
    }
#endif
    std::cout << "=== QUIC Version Negotiation Test ===" << std::endl;

    const int kQuicPort = 19548;
    int rc = 0;
    std::thread quicServerThread;

    try {
        x509cert cert;
        if (!cert.loadFromBuffer(test_cert_der)) {
            throw std::runtime_error("failed to load certificate");
        }
        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;

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

        // --- Craft a minimal long-header Initial with an unsupported version ---
        std::cout << "\n--- Test: unsupported version triggers a Version Negotiation reply ---" << std::endl;

        std::vector<uint8_t> client_dcid(8), client_scid(8);
        fillRandomBytes(client_dcid.data(), client_dcid.size());
        fillRandomBytes(client_scid.data(), client_scid.size());
        const uint32_t kUnsupportedVersion = 0xdeadbeef;

        std::vector<uint8_t> pkt;
        pkt.push_back(0xC0); // long header, fixed bit, type=Initial(00)
        pkt.push_back(static_cast<uint8_t>(kUnsupportedVersion >> 24));
        pkt.push_back(static_cast<uint8_t>(kUnsupportedVersion >> 16));
        pkt.push_back(static_cast<uint8_t>(kUnsupportedVersion >> 8));
        pkt.push_back(static_cast<uint8_t>(kUnsupportedVersion));
        pkt.push_back(static_cast<uint8_t>(client_dcid.size()));
        pkt.insert(pkt.end(), client_dcid.begin(), client_dcid.end());
        pkt.push_back(static_cast<uint8_t>(client_scid.size()));
        pkt.insert(pkt.end(), client_scid.begin(), client_scid.end());

        native_socket_t rawfd = ::socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        check(rawfd != kInvalidSocket, "raw UDP socket created");
        if (rawfd == kInvalidSocket) {
#ifdef _WIN32
            std::cerr << "socket() failed: " << ::WSAGetLastError() << std::endl;
#else
            std::cerr << "socket() failed" << std::endl;
#endif
            throw std::runtime_error("unable to create raw UDP socket");
        }

        sockaddr_in serverAddr{};
        serverAddr.sin_family = AF_INET;
        serverAddr.sin_port = htons(kQuicPort);
        ::inet_pton(AF_INET, "127.0.0.1", &serverAddr.sin_addr);

#ifdef _WIN32
        const int sent = ::sendto(rawfd,
                                  reinterpret_cast<const char*>(pkt.data()),
                                  static_cast<int>(pkt.size()), 0,
                                  reinterpret_cast<const sockaddr*>(&serverAddr),
                                  static_cast<int>(sizeof(serverAddr)));
#else
        const ssize_t sent = ::sendto(rawfd, pkt.data(), pkt.size(), 0,
                                      reinterpret_cast<const sockaddr*>(&serverAddr),
                                      sizeof(serverAddr));
#endif
        check(sent == static_cast<decltype(sent)>(pkt.size()), "crafted Initial sent to server");

        std::vector<uint8_t> resp(2048);
        const bool readable = waitForReadable(rawfd, 5000);
        check(readable, "received a reply within 5s (not silence)");

        if (readable) {
#ifdef _WIN32
            const int n = ::recv(rawfd, reinterpret_cast<char*>(resp.data()),
                                 static_cast<int>(resp.size()), 0);
#else
            const ssize_t n = ::recv(rawfd, resp.data(), resp.size(), 0);
#endif
            check(n > 0, "reply has non-zero length");
            resp.resize(n > 0 ? static_cast<size_t>(n) : 0);

            // Header form bit must be set (long header).
            check(!resp.empty() && (resp[0] & 0x80) != 0, "reply is a long-header packet");

            // Version field must be exactly 0 (the VN marker).
            bool versionIsZero = resp.size() >= 5 &&
                resp[1] == 0 && resp[2] == 0 && resp[3] == 0 && resp[4] == 0;
            check(versionIsZero, "reply version field == 0 (Version Negotiation marker)");

            // DCID must equal our SCID, SCID must equal our DCID (swapped echo).
            size_t off = 5;
            bool dcidOk = false, scidOk = false;
            std::vector<uint32_t> versions;
            if (off < resp.size()) {
                uint8_t dcid_len = resp[off++];
                if (off + dcid_len <= resp.size()) {
                    std::vector<uint8_t> got_dcid(resp.begin() + off, resp.begin() + off + dcid_len);
                    dcidOk = (got_dcid == client_scid);
                    off += dcid_len;

                    if (off < resp.size()) {
                        uint8_t scid_len = resp[off++];
                        if (off + scid_len <= resp.size()) {
                            std::vector<uint8_t> got_scid(resp.begin() + off, resp.begin() + off + scid_len);
                            scidOk = (got_scid == client_dcid);
                            off += scid_len;

                            while (off + 4 <= resp.size()) {
                                uint32_t v = (static_cast<uint32_t>(resp[off]) << 24) |
                                             (static_cast<uint32_t>(resp[off + 1]) << 16) |
                                             (static_cast<uint32_t>(resp[off + 2]) << 8) |
                                             static_cast<uint32_t>(resp[off + 3]);
                                versions.push_back(v);
                                off += 4;
                            }
                        }
                    }
                }
            }
            check(dcidOk, "reply DCID == our original SCID (echoed, swapped)");
            check(scidOk, "reply SCID == our original DCID (echoed, swapped)");

            bool hasV1 = false, hasV2 = false;
            for (auto v : versions) {
                if (v == quic::QUIC_VERSION_1) hasV1 = true;
                if (v == quic::QUIC_VERSION_2) hasV2 = true;
            }
            check(hasV1, "supported version list includes QUICv1 (0x00000001)");
            check(hasV2, "supported version list includes QUICv2 (0x6b3343cf)");
        }

        closeNativeSocket(rawfd);

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

    if (quicServerThread.joinable()) {
        event::Running = false;
        quicServerThread.join();
    }

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

    const int result = (rc != 0 || g_failed > 0) ? 1 : 0;
#ifdef _WIN32
    ::WSACleanup();
#endif
    return result;
}