#include <arpa/inet.h>
#include <netdb.h>
#include <sys/socket.h>
#include <unistd.h>
#include <cstring>
#include <iostream>
#include <random>
#include <string>
#include <vector>

static std::string base64(const unsigned char* data, size_t len) {
    static const char* tbl = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::string out;
    out.reserve(((len + 2) / 3) * 4);
    size_t i = 0;
    while (i + 2 < len) {
        unsigned int n = (data[i] << 16) | (data[i+1] << 8) | data[i+2];
        out.push_back(tbl[(n >> 18) & 63]);
        out.push_back(tbl[(n >> 12) & 63]);
        out.push_back(tbl[(n >> 6) & 63]);
        out.push_back(tbl[n & 63]);
        i += 3;
    }
    if (i < len) {
        unsigned int n = (data[i] << 16);
        if (i + 1 < len) n |= (data[i+1] << 8);
        out.push_back(tbl[(n >> 18) & 63]);
        out.push_back(tbl[(n >> 12) & 63]);
        if (i + 1 < len) {
            out.push_back(tbl[(n >> 6) & 63]);
            out.push_back('=');
        } else {
            out.push_back('=');
            out.push_back('=');
        }
    }
    return out;
}

static int connect_tcp(const std::string& host, int port) {
    int sock = ::socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) return -1;
    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    if (::inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) {
        ::close(sock);
        return -1;
    }
    if (::connect(sock, (sockaddr*)&addr, sizeof(addr)) != 0) {
        ::close(sock);
        return -1;
    }
    return sock;
}

static bool send_all(int sock, const void* buf, size_t len) {
    const char* p = (const char*)buf;
    while (len) {
        ssize_t n = ::send(sock, p, len, 0);
        if (n <= 0) return false;
        p += n; len -= (size_t)n;
    }
    return true;
}

static std::string recv_headers(int sock) {
    std::string h;
    char c;
    while (true) {
        ssize_t n = ::recv(sock, &c, 1, 0);
        if (n <= 0) break;
        h.push_back(c);
        if (h.size() >= 4 && h.substr(h.size()-4) == "\r\n\r\n") break;
    }
    return h;
}

static bool send_ws_text(int sock, const std::string& payload) {
    std::vector<unsigned char> frame;
    frame.reserve(2 + 4 + payload.size());
    unsigned char b0 = 0x80 | 0x1; // FIN + text
    frame.push_back(b0);
    size_t len = payload.size();
    if (len < 126) {
        frame.push_back(0x80 | (unsigned char)len); // mask bit set
    } else if (len <= 0xFFFF) {
        frame.push_back(0x80 | 126);
        frame.push_back((len >> 8) & 0xFF);
        frame.push_back(len & 0xFF);
    } else {
        frame.push_back(0x80 | 127);
        for (int i=7;i>=0;--i) frame.push_back((len >> (i*8)) & 0xFF);
    }
    unsigned char mask[4];
    std::random_device rd; std::mt19937 rng(rd()); std::uniform_int_distribution<int> dist(0,255);
    for (int i=0;i<4;++i) mask[i] = (unsigned char)dist(rng);
    frame.insert(frame.end(), mask, mask+4);
    for (size_t i=0;i<len;++i) frame.push_back((unsigned char)(payload[i] ^ mask[i%4]));
    return send_all(sock, frame.data(), frame.size());
}

static std::string recv_ws_text(int sock) {
    unsigned char hdr[2];
    if (::recv(sock, hdr, 2, MSG_WAITALL) != 2) return {};
    bool fin = (hdr[0] & 0x80) != 0; (void)fin;
    unsigned char opcode = hdr[0] & 0x0F; (void)opcode;
    bool masked = (hdr[1] & 0x80) != 0;
    uint64_t len = hdr[1] & 0x7F;
    if (len == 126) {
        unsigned char ext[2];
        if (::recv(sock, ext, 2, MSG_WAITALL) != 2) return {};
        len = (ext[0] << 8) | ext[1];
    } else if (len == 127) {
        unsigned char ext[8];
        if (::recv(sock, ext, 8, MSG_WAITALL) != 8) return {};
        len = 0; for (int i=0;i<8;++i) len = (len<<8) | ext[i];
    }
    unsigned char mask[4] = {0,0,0,0};
    if (masked) {
        if (::recv(sock, mask, 4, MSG_WAITALL) != 4) return {};
    }
    std::string payload; payload.resize((size_t)len);
    if (len > 0) {
        if (::recv(sock, payload.data(), len, MSG_WAITALL) != (ssize_t)len) return {};
        if (masked) {
            for (size_t i=0;i<len;++i) payload[i] ^= mask[i%4];
        }
    }
    return payload;
}

int main() {
    int sock = connect_tcp("127.0.0.1", 1235);
    if (sock < 0) {
        std::cerr << "connect failed" << std::endl; return 1;
    }
    unsigned char keyraw[16];
    std::random_device rd; for (int i=0;i<16;++i) keyraw[i] = (unsigned char)rd();
    std::string key = base64(keyraw, 16);
    std::string req;
    req += "GET / HTTP/1.1\r\n";
    req += "Host: localhost:1235\r\n";
    req += "Upgrade: websocket\r\n";
    req += "Connection: Upgrade\r\n";
    req += "Sec-WebSocket-Key: " + key + "\r\n";
    req += "Sec-WebSocket-Version: 13\r\n\r\n";
    if (!send_all(sock, req.c_str(), req.size())) { std::cerr << "send handshake failed" << std::endl; return 2; }
    std::string resp = recv_headers(sock);
    std::cout << "Handshake response:\n" << resp << std::endl;
    std::string init = R"({"type":"init","name":"Probe","room":"Lobby","color":"#000000","password":""})";
    if (!send_ws_text(sock, init)) { std::cerr << "send frame failed" << std::endl; return 3; }
    std::string msg = recv_ws_text(sock);
    std::cout << "First WS message from server: " << msg << std::endl;
    ::close(sock);
    return 0;
}