#include "rpc/ws_frame.hpp"
#include "rpc/ws_handshake.hpp"

#include <string>
#include <vector>

#include "check.hpp"

using namespace velox::daemon::rpc;

namespace {

// Build a *client* frame: FIN/opcode, mask bit set, a fixed 4-byte mask, masked payload.
std::string client_frame(WsOpcode op, std::string_view payload, bool fin = true) {
    std::string f;
    f.push_back(static_cast<char>((fin ? 0x80 : 0x00) | static_cast<std::uint8_t>(op)));
    const std::size_t n = payload.size();
    if (n < 126) {
        f.push_back(static_cast<char>(0x80 | n));
    } else if (n <= 0xFFFF) {
        f.push_back(static_cast<char>(0x80 | 126));
        f.push_back(static_cast<char>((n >> 8) & 0xFF));
        f.push_back(static_cast<char>(n & 0xFF));
    } else {
        f.push_back(static_cast<char>(0x80 | 127));
        for (int i = 7; i >= 0; --i)
            f.push_back(static_cast<char>((static_cast<std::uint64_t>(n) >> (i * 8)) & 0xFF));
    }
    const char key[4] = {0x12, 0x34, 0x56, 0x78};
    f.append(key, 4);
    for (std::size_t i = 0; i < n; ++i) f.push_back(static_cast<char>(payload[i] ^ key[i & 3]));
    return f;
}

}  // namespace

void run() {
    // --- one text frame ------------------------------------------------------------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        CHECK(r.feed(client_frame(WsOpcode::Text, "{\"a\":1}"), m) == WsFrameReader::Status::Ok);
        CHECK_EQ(m.size(), 1u);
        CHECK(m[0].opcode == WsOpcode::Text);
        CHECK_EQ(m[0].payload, std::string("{\"a\":1}"));
    }

    // --- fragmented: text (fin=0) + continuation (fin=1) --------------------------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        r.feed(client_frame(WsOpcode::Text, "hel", /*fin=*/false), m);
        CHECK_EQ(m.size(), 0u);
        r.feed(client_frame(WsOpcode::Continuation, "lo", /*fin=*/true), m);
        CHECK_EQ(m.size(), 1u);
        CHECK_EQ(m[0].payload, std::string("hello"));
    }

    // --- byte-at-a-time delivery still reassembles -------------------------------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        const std::string frame = client_frame(WsOpcode::Text, "streamed");
        for (char ch : frame) r.feed(std::string_view(&ch, 1), m);
        CHECK_EQ(m.size(), 1u);
        CHECK_EQ(m[0].payload, std::string("streamed"));
    }

    // --- a 200-byte payload exercises the 16-bit length path --------------------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        const std::string big(200, 'x');
        r.feed(client_frame(WsOpcode::Text, big), m);
        CHECK_EQ(m.size(), 1u);
        CHECK_EQ(m[0].payload.size(), 200u);
    }

    // --- ping is surfaced so the server can pong -------------------------------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        r.feed(client_frame(WsOpcode::Ping, "hi"), m);
        CHECK_EQ(m.size(), 1u);
        CHECK(m[0].opcode == WsOpcode::Ping);
    }

    // --- an unmasked client frame is a protocol error (RFC 6455 §5.1) ----------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        std::string bad;
        bad.push_back(static_cast<char>(0x81));  // FIN + text
        bad.push_back(static_cast<char>(0x03));  // len 3, mask bit clear
        bad.append("abc");
        CHECK(r.feed(bad, m) == WsFrameReader::Status::ProtocolError);
    }

    // --- a declared length past the cap is rejected before allocating ----------
    {
        WsFrameReader r;
        std::vector<WsMessage> m;
        std::string hdr;
        hdr.push_back(static_cast<char>(0x82));  // FIN + binary
        hdr.push_back(static_cast<char>(0x80 | 127));
        for (int i = 7; i >= 0; --i)
            hdr.push_back(static_cast<char>((0x0000000001000000ull >> (i * 8)) & 0xFF));  // 16 MiB
        CHECK(r.feed(hdr, m) == WsFrameReader::Status::MessageTooBig);
    }

    // --- ws_encode: server frames are unmasked, correct length byte ------------
    {
        const std::string f = ws_encode(WsOpcode::Text, "abc");
        CHECK_EQ(static_cast<std::uint8_t>(f[0]), 0x81u);
        CHECK_EQ(static_cast<std::uint8_t>(f[1]), 0x03u);  // len 3, no mask bit
        CHECK_EQ(f.substr(2), std::string("abc"));
    }

    // --- RFC 6455 §1.3 sample accept value ------------------------------------
    CHECK_EQ(ws_accept_key("dGhlIHNhbXBsZSBub25jZQ=="),
             std::string("s3pPLMBiTxaQ9kYGzzhZRbK+xOo="));

    // --- handshake: a page origin is refused, an extension origin upgrades -----
    {
        const std::string req_page =
            "GET / HTTP/1.1\r\nHost: 127.0.0.1:52000\r\nUpgrade: websocket\r\n"
            "Connection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
            "Sec-WebSocket-Version: 13\r\nOrigin: https://evil.example\r\n\r\n";
        const auto r = ws_try_handshake(req_page);
        CHECK(r.complete);
        CHECK(!r.ok);
        CHECK(r.response.find("403") != std::string::npos);
    }
    {
        const std::string req_ext =
            "GET / HTTP/1.1\r\nHost: 127.0.0.1:52000\r\nUpgrade: websocket\r\n"
            "Connection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
            "Sec-WebSocket-Version: 13\r\n"
            "Origin: moz-extension://11111111-2222-3333-4444-555555555555\r\n\r\n";
        const auto r = ws_try_handshake(req_ext);
        CHECK(r.complete);
        CHECK(r.ok);
        CHECK(r.response.find("101") != std::string::npos);
        CHECK(r.response.find("Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=") !=
              std::string::npos);
        CHECK_EQ(r.origin, std::string("moz-extension://11111111-2222-3333-4444-555555555555"));
    }
}

TEST_MAIN()