// RFC 9000 (QUIC Transport) compliance unit tests for libnetplus
// Tests transport-layer semantics: flow control, stream types,
// connection state, data integrity, and congestion behavior.

#include <iostream>
#include <string>
#include <cassert>
#include <cstring>
#include <memory>
#include <vector>
#include <thread>
#include <atomic>
#include <chrono>
#include <numeric>
#include <algorithm>

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

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

using namespace netplus;

static std::atomic<bool> g_server_ready(false);
static int g_passed = 0, g_failed = 0;

// ============================================================================
// Helpers
// ============================================================================

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

// Fill a buffer with a deterministic pattern for verification
static void fillPattern(std::vector<uint8_t>& buf, uint8_t seed = 0) {
    for (size_t i = 0; i < buf.size(); i++)
        buf[i] = static_cast<uint8_t>((i + seed) & 0xFF);
}

// Receive all stream data until stream is complete or deadline expires
static size_t drainStream(quic& client, uint64_t sid, std::vector<uint8_t>& out,
                          int timeout_sec = 10) {
    size_t total = 0;
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_sec);
    while (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; }
        }
        if (client.isStreamComplete(sid)) break;
        std::this_thread::sleep_for(std::chrono::microseconds(200));
    }
    return total;
}

// Connect a QUIC client to the test server
static void connectClient(quic& client) {
    client.connect("127.0.0.1", 18443);
}

// ============================================================================
// Echo server (receives stream data, echoes it back)
// ============================================================================

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

    void RequestEvent(con& curcon, const int tid, ULONG_PTR args) override {
        if (!curcon.RecvData.empty()) {
            curcon.SendData.append(curcon.RecvData.data(), curcon.RecvData.size());
            curcon.RecvData.clear();
        }
    }
    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);
    }
};

// ============================================================================
// Server thread
// ============================================================================

static void runServer(std::map<std::string, ssl::CertificateBundle>& certs) {
    try {
        quic serverSock(certs, "127.0.0.1", 18443, 1024, -1);
        // Raised so test_pmtu_discovery()'s client can search past the
        // 1472 baseline; clients that never raise their own ceiling still
        // clamp their search to 1472 regardless (see setPmtuCeiling()), so
        // this has no effect on the other tests in this file.
        serverSock.setPmtuCeiling(9000);
        serverSock.setStreamCallback([](netplus::socket* sock, uint64_t stream_id,
                                        const std::vector<uint8_t>& data, bool fin) {
            quic* q = dynamic_cast<quic*>(sock);
            if (!q || data.empty()) return;
            q->sendStreamData(stream_id, data, fin);
        });

        g_server_ready.store(true);
        EchoServer srv({&serverSock});
        srv.runEventloop();
    } catch (NetException& e) {
        std::cerr << "[Server] NetException: " << e.what() << std::endl;
    } catch (std::exception& e) {
        std::cerr << "[Server] Exception: " << e.what() << std::endl;
    }
}

// Wait for server to be ready
static void waitServer() {
    while (!g_server_ready.load())
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    std::this_thread::sleep_for(std::chrono::milliseconds(200));
}

// ============================================================================
// RFC 9000 §5: Connection State Machine
// ============================================================================

static void test_connection_state() {
    std::cout << "\n=== RFC 9000 §5: Connection State ===" << std::endl;

    // §5.1: Before connect, state should be Idle
    {
        quic client;
        check(client.getConnectionState() == quic::ConnectionState::Idle,
              "Initial state is Idle");
    }

    // §5.1: After successful handshake, state should be Connected
    {
        quic client;
        connectClient(client);
        check(client.getConnectionState() == quic::ConnectionState::Connected,
              "After connect, state is Connected");
        check(client.getHandshakeDone(), "Handshake marked complete");
    }

    // §5.2: After close, state should transition to Draining or Closed
    {
        quic client;
        connectClient(client);
        client.close();
        auto st = client.getConnectionState();
        check(st == quic::ConnectionState::Draining ||
              st == quic::ConnectionState::Closed,
              "After close, state is Draining or Closed");
    }
}

// ============================================================================
// RFC 9000 §4: Flow Control — Connection-Level
// ============================================================================

static void test_connection_flow_control() {
    std::cout << "\n=== RFC 9000 §4: Connection-Level Flow Control ===" << std::endl;

    // Send data that fits within default flow control window — should succeed
    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        std::vector<uint8_t> payload(32768); // 32 KB
        fillPattern(payload, 0x11);

        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(payload.size() + 4096);
        size_t total = drainStream(client, sid, recv_buf, 10);

        check(total == payload.size(), "32KB echo: correct size returned");
        check(std::memcmp(payload.data(), recv_buf.data(), payload.size()) == 0,
              "32KB echo: data integrity");
    }
}

// ============================================================================
// RFC 9000 §4.1: Stream-Level Flow Control
// ============================================================================

static void test_stream_flow_control() {
    std::cout << "\n=== RFC 9000 §4.1: Stream-Level Flow Control ===" << std::endl;

    // Multiple streams on same connection — each gets independent flow control
    {
        quic client;
        connectClient(client);

        uint64_t s1 = client.openStream(true);
        uint64_t s2 = client.openStream(true);

        check(s1 != s2, "Two streams have different IDs");

        std::vector<uint8_t> p1(4096, 0xAA);
        std::vector<uint8_t> p2(8192, 0xBB);

        client.sendStreamData(s1, p1, true);
        client.sendStreamData(s2, p2, true);

        std::vector<uint8_t> r1(p1.size() + 1024), r2(p2.size() + 1024);
        size_t t1 = drainStream(client, s1, r1, 10);
        size_t t2 = drainStream(client, s2, r2, 10);

        check(t1 == p1.size(), "Stream 1: correct echo size");
        check(t2 == p2.size(), "Stream 2: correct echo size");
        check(std::memcmp(p1.data(), r1.data(), p1.size()) == 0,
              "Stream 1: data integrity");
        check(std::memcmp(p2.data(), r2.data(), p2.size()) == 0,
              "Stream 2: data integrity");
    }
}

// ============================================================================
// RFC 9000 §2.1: Stream Types — Bidirectional
// ============================================================================

static void test_bidi_streams() {
    std::cout << "\n=== RFC 9000 §2.1: Bidirectional Streams ===" << std::endl;

    // Client-initiated bidi stream IDs are 0, 4, 8, ... (stream_id % 4 == 0)
    {
        quic client;
        connectClient(client);

        uint64_t s1 = client.openStream(true);
        uint64_t s2 = client.openStream(true);
        uint64_t s3 = client.openStream(true);

        check(s1 % 4 == 0, "First bidi stream ID divisible by 4");
        check(s2 % 4 == 0, "Second bidi stream ID divisible by 4");
        check(s3 % 4 == 0, "Third bidi stream ID divisible by 4");
        check(s1 < s2 && s2 < s3, "Stream IDs are monotonically increasing");
    }
}

// ============================================================================
// RFC 9000 §2.2: Stream Types — Unidirectional
// ============================================================================

static void test_uni_streams() {
    std::cout << "\n=== RFC 9000 §2.2: Unidirectional Streams ===" << std::endl;

    // Client-initiated uni stream IDs are 2, 6, 10, ... (stream_id % 4 == 2)
    {
        quic client;
        connectClient(client);

        uint64_t s1 = client.openStream(false);
        uint64_t s2 = client.openStream(false);

        check(s1 % 4 == 2, "First uni stream ID % 4 == 2");
        check(s2 % 4 == 2, "Second uni stream ID % 4 == 2");
        check(s1 < s2, "Uni stream IDs are increasing");
    }
}

// ============================================================================
// RFC 9000 §3.2: Send Stream States — FIN handling
// ============================================================================

static void test_fin_handling() {
    std::cout << "\n=== RFC 9000 §3.2: FIN Handling ===" << std::endl;

    // Send data with FIN in a single call
    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        std::vector<uint8_t> payload(512, 0xCC);
        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(payload.size() + 1024);
        size_t total = drainStream(client, sid, recv_buf, 5);
        check(total == payload.size(), "FIN in single send: echoed correctly");
    }

    // Send data first, then FIN separately (empty data + fin)
    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        std::vector<uint8_t> payload(512, 0xDD);
        client.sendStreamData(sid, payload.data(), payload.size(), false);
        // Send FIN with empty data
        client.sendStreamData(sid, nullptr, 0, true);

        std::vector<uint8_t> recv_buf(payload.size() + 1024);
        size_t total = drainStream(client, sid, recv_buf, 5);
        check(total == payload.size(), "Separate FIN: echoed correctly");
    }
}

// ============================================================================
// RFC 9000 §7.4: Transport Parameters — Version Negotiation
// ============================================================================

static void test_version() {
    std::cout << "\n=== RFC 9000 §7.4: Version ===" << std::endl;

    {
        quic client;
        connectClient(client);
        uint32_t ver = client.getVersion();
        check(ver == quic::QUIC_VERSION_1 || ver == quic::QUIC_VERSION_2,
              "Negotiated version is QUIC v1 or v2");
    }
}

// ============================================================================
// RFC 9000 §5.1: Connection IDs
// ============================================================================

static void test_connection_ids() {
    std::cout << "\n=== RFC 9000 §5.1: Connection IDs ===" << std::endl;

    {
        quic client;
        connectClient(client);

        auto& local = client.getLocalConnectionId();
        auto& remote = client.getRemoteConnectionId();

        check(!local.empty(), "Local CID is non-empty");
        check(!remote.empty(), "Remote CID is non-empty");
        check(local != remote, "Local and remote CIDs are different");
        check(local.size() <= 20, "Local CID <= 20 bytes (RFC 9000 §17.2)");
        check(remote.size() <= 20, "Remote CID <= 20 bytes");
    }
}

// ============================================================================
// RFC 9000 §10: Data Integrity — Small payloads
// ============================================================================

static void test_data_integrity_small() {
    std::cout << "\n=== RFC 9000 §10: Data Integrity (Small) ===" << std::endl;

    const size_t sizes[] = {1, 16, 127, 256, 1000, 1200};
    for (size_t sz : sizes) {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        std::vector<uint8_t> payload(sz);
        fillPattern(payload, static_cast<uint8_t>(sz));

        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(sz + 4096);
        size_t total = drainStream(client, sid, recv_buf, 5);

        std::string label = std::to_string(sz) + "-byte echo integrity";
        check(total == sz && std::memcmp(payload.data(), recv_buf.data(), sz) == 0,
              label.c_str());
    }
}

// ============================================================================
// RFC 9000 §10: Data Integrity — Large payloads (multi-packet)
// ============================================================================

static void test_data_integrity_large() {
    std::cout << "\n=== RFC 9000 §10: Data Integrity (Large) ===" << std::endl;

    // 256 KB — requires many QUIC packets, exercises reassembly and flow control
    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        const size_t sz = 256 * 1024;
        std::vector<uint8_t> payload(sz);
        fillPattern(payload, 0x42);

        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(sz + 4096);
        size_t total = drainStream(client, sid, recv_buf, 15);

        check(total == sz, "256KB echo: correct size");
        check(std::memcmp(payload.data(), recv_buf.data(), sz) == 0,
              "256KB echo: data integrity");
    }

    // 1 MB — exercises congestion control and MAX_DATA updates
    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        const size_t sz = 1024 * 1024;
        std::vector<uint8_t> payload(sz);
        fillPattern(payload, 0x77);

        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(sz + 4096);
        size_t total = drainStream(client, sid, recv_buf, 20);

        check(total == sz, "1MB echo: correct size");
        check(std::memcmp(payload.data(), recv_buf.data(), sz) == 0,
              "1MB echo: data integrity");
    }
}

// ============================================================================
// RFC 9000 §2.4: Multiple Concurrent Streams
// ============================================================================

static void test_concurrent_streams() {
    std::cout << "\n=== RFC 9000 §2.4: Concurrent Streams ===" << std::endl;

    {
        quic client;
        connectClient(client);

        const int N = 4;
        std::vector<uint64_t> sids(N);
        std::vector<std::vector<uint8_t>> payloads(N);

        for (int i = 0; i < N; i++) {
            sids[i] = client.openStream(true);
            payloads[i].resize(4096 * (i + 1));
            fillPattern(payloads[i], static_cast<uint8_t>(i * 37));
            client.sendStreamData(sids[i], payloads[i], true);
        }

        bool all_ok = true;
        for (int i = 0; i < N; i++) {
            std::vector<uint8_t> recv_buf(payloads[i].size() + 4096);
            size_t total = drainStream(client, sids[i], recv_buf, 15);
            if (total != payloads[i].size() ||
                std::memcmp(payloads[i].data(), recv_buf.data(), payloads[i].size()) != 0) {
                all_ok = false;
                std::cout << "    Stream " << sids[i] << ": expected " << payloads[i].size()
                          << " got " << total << std::endl;
            }
        }
        check(all_ok, "4 concurrent streams echoed correctly");
    }
}

// ============================================================================
// RFC 9000 §19.8: STREAM frame with zero-length data + FIN
// ============================================================================

static void test_empty_fin() {
    std::cout << "\n=== RFC 9000 §19.8: Empty FIN ===" << std::endl;

    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        // Send data, then close with empty FIN
        std::vector<uint8_t> payload(128, 0xEE);
        size_t sent = client.sendStreamData(sid, payload.data(), payload.size(), false);
        check(sent == payload.size(), "Initial data sent successfully");

        // Empty data + FIN
        std::vector<uint8_t> empty;
        client.sendStreamData(sid, empty, true);

        std::vector<uint8_t> recv_buf(payload.size() + 4096);
        size_t total = drainStream(client, sid, recv_buf, 5);
        check(total == payload.size(), "Empty FIN: echo size correct");
    }
}

// ============================================================================
// RFC 9000 §4.6: MAX_STREAMS — Rapid stream creation
// ============================================================================

static void test_rapid_stream_creation() {
    std::cout << "\n=== RFC 9000 §4.6: Rapid Stream Creation ===" << std::endl;

    {
        quic client;
        connectClient(client);

        const int N = 16;
        bool all_ok = true;

        for (int i = 0; i < N; i++) {
            uint64_t sid = client.openStream(true);
            std::vector<uint8_t> payload(64, static_cast<uint8_t>(i));
            client.sendStreamData(sid, payload, true);

            std::vector<uint8_t> recv_buf(payload.size() + 1024);
            size_t total = drainStream(client, sid, recv_buf, 3);
            if (total != payload.size() ||
                std::memcmp(payload.data(), recv_buf.data(), payload.size()) != 0) {
                all_ok = false;
                std::cout << "    Stream " << i << " (id=" << sid << "): failed" << std::endl;
            }
        }
        check(all_ok, "16 sequential streams echoed correctly");
    }
}

// ============================================================================
// RFC 9000 §17.2: Handshake — TLS 1.3 Cipher Selection
// ============================================================================

static void test_cipher_selection() {
    std::cout << "\n=== RFC 9000 §17.2: Cipher Selection ===" << std::endl;

    {
        quic client;
        connectClient(client);
        uint16_t cipher = client.getSelectedCipher();
        // RFC 9001 §5.2 mandates TLS_AES_128_GCM_SHA256 (0x1301)
        // Also allows TLS_AES_256_GCM_SHA384 (0x1302) or TLS_CHACHA20_POLY1305_SHA256 (0x1303)
        check(cipher == 0x1301 || cipher == 0x1302 || cipher == 0x1303,
              "Selected cipher is a valid TLS 1.3 AEAD");
        std::cout << "    Cipher: 0x" << std::hex << cipher << std::dec << std::endl;
    }
}

// ============================================================================
// RFC 9000 §10.1: Idle Timeout — Connection survives short idle
// ============================================================================

static void test_idle_survival() {
    std::cout << "\n=== RFC 9000 §10.1: Idle Survival ===" << std::endl;

    {
        quic client;
        connectClient(client);
        uint64_t sid = client.openStream(true);

        // Send first chunk
        std::vector<uint8_t> p1(256, 0x11);
        client.sendStreamData(sid, p1.data(), p1.size(), false);

        // Wait 2 seconds (connection should survive via PING keep-alive)
        std::this_thread::sleep_for(std::chrono::seconds(2));

        // Send second chunk + FIN
        std::vector<uint8_t> p2(256, 0x22);
        client.sendStreamData(sid, p2.data(), p2.size(), true);

        // Expected echo: p1 + p2
        std::vector<uint8_t> expected;
        expected.insert(expected.end(), p1.begin(), p1.end());
        expected.insert(expected.end(), p2.begin(), p2.end());

        std::vector<uint8_t> recv_buf(expected.size() + 4096);
        size_t total = drainStream(client, sid, recv_buf, 10);

        check(total == expected.size(), "Idle survival: correct total size");
        check(std::memcmp(expected.data(), recv_buf.data(), expected.size()) == 0,
              "Idle survival: data integrity after 2s idle");
    }
}

// ============================================================================
// RFC 9000 §14.3-14.4 / RFC 8899: Path MTU Discovery
// ============================================================================

static void test_pmtu_discovery() {
    std::cout << "\n=== RFC 9000 §14.3-14.4: Path MTU Discovery ===" << std::endl;

    // Default ceiling (1472) must be a no-op: discovery never raises the
    // confirmed send size past today's fixed baseline unless asked to.
    {
        quic client;
        connectClient(client);
        auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
        while (std::chrono::steady_clock::now() < deadline) {
            client.pumpNetwork(MSG_DONTWAIT);
            std::this_thread::sleep_for(std::chrono::microseconds(200));
        }
        check(client.getMaxUdpPayload() == 1472,
              "Default ceiling: confirmed size stays at the 1472 baseline");
    }

    // Raising the ceiling should let the client discover a larger usable
    // size over the loopback path (server's ceiling raised in runServer()).
    {
        quic client;
        client.setPmtuCeiling(9000);
        connectClient(client);

        auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
        while (std::chrono::steady_clock::now() < deadline) {
            client.pumpNetwork(MSG_DONTWAIT);
            std::this_thread::sleep_for(std::chrono::microseconds(200));
        }

        uint64_t discovered = client.getMaxUdpPayload();
        std::cout << "    Discovered max_udp_payload: " << discovered << std::endl;
        check(discovered > 1472, "Raised ceiling: confirmed size grows past the 1472 baseline");

        // Stream sends must still round-trip correctly once chunked at
        // the larger discovered datagram size.
        uint64_t sid = client.openStream(true);
        std::vector<uint8_t> payload(16384);
        fillPattern(payload, 0x33);
        client.sendStreamData(sid, payload, true);

        std::vector<uint8_t> recv_buf(payload.size() + 4096);
        size_t total = drainStream(client, sid, recv_buf, 10);
        check(total == payload.size(), "Post-discovery echo: correct size returned");
        check(std::memcmp(payload.data(), recv_buf.data(), payload.size()) == 0,
              "Post-discovery echo: data integrity at larger datagram size");
    }
}

// ============================================================================
// Main
// ============================================================================

int main() {
    std::cout << "=== RFC 9000 QUIC Transport Compliance Tests ===" << std::endl;

    // Load certificates
    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;

    // Start echo server
    std::thread server_thread(runServer, std::ref(certs));
    waitServer();

    // Run all tests
    test_connection_state();
    test_version();
    test_connection_ids();
    test_cipher_selection();
    test_bidi_streams();
    test_uni_streams();
    test_fin_handling();
    test_empty_fin();
    test_data_integrity_small();
    test_data_integrity_large();
    test_connection_flow_control();
    test_stream_flow_control();
    test_concurrent_streams();
    test_rapid_stream_creation();
    test_idle_survival();
    test_pmtu_discovery();

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

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

    return g_failed > 0 ? 1 : 0;
}
