// g++ -o websocket websocket.cpp -lssl -lcrypto
// js
// const ws = new WebSocket('ws:localhost:8080')
// ws.onmessage = e => console.log(e.data)
// ws.send('hello')

#include <iostream>
#include <string>
#include <cstring>
#include <vector>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sstream>
#include <iomanip>
#include <openssl/sha.h>
#include <openssl/evp.h>

using namespace std;

// Base64 кодирование (упрощенная версия)
string base64_encode(const string &in) {
    const string base64_chars =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
        "abcdefghijklmnopqrstuvwxyz"
        "0123456789+/";

    string out;
    int val = 0, valb = -6;

    for (unsigned char c : in) {
        val = (val << 8) + c;
        valb += 8;
        while (valb >= 0) {
            out.push_back(base64_chars[(val >> valb) & 0x3F]);
            valb -= 6;
        }
    }

    if (valb > -6)
        out.push_back(base64_chars[((val << 8) >> (valb + 8)) & 0x3F]);

    while (out.size() % 4)
        out.push_back('=');

    return out;
}

// Генерация ответа для WebSocket handshake
string get_websocket_accept(const string &client_key) {
    const string magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    string combined = client_key + magic;

    unsigned char digest[SHA_DIGEST_LENGTH];
    SHA1(reinterpret_cast<const unsigned char*>(combined.c_str()), combined.size(), digest);

    return base64_encode(string(reinterpret_cast<char*>(digest), SHA_DIGEST_LENGTH));
}

// Демаскировка данных
void unmask_data(char *data, size_t len, const char *mask) {
    for(size_t i = 0; i < len; i++) {
        data[i] ^= mask[i % 4];
    }
}

int main() {
    const int port = 8080;

    // Создание сокета
    int server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd < 0) {
        cerr << "Ошибка создания сокета" << endl;
        return 1;
    }

    // Настройка адреса
    sockaddr_in address;
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(port);

    // Бинд и листен
    if (bind(server_fd, (sockaddr*)&address, sizeof(address)) < 0) {
        cerr << "Ошибка бинда" << endl;
        return 1;
    }

    if (listen(server_fd, 1) < 0) {
        cerr << "Ошибка листена" << endl;
        return 1;
    }

    cout << "Сервер запущен на порту " << port << endl;

    while(true) {
        // Принятие соединения
        sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);
        int client_socket = accept(server_fd, (sockaddr*)&client_addr, &client_len);

        if (client_socket < 0) {
            cerr << "Ошибка принятия соединения" << endl;
            continue;
        }

        // Чтение HTTP-запроса
        char buffer[4096] = {0};
        read(client_socket, buffer, 4096);

        // Поиск ключа WebSocket
        string request(buffer);
        size_t key_pos = request.find("Sec-WebSocket-Key: ");
        if(key_pos == string::npos) {
            cerr << "Неверный WebSocket запрос" << endl;
            close(client_socket);
            continue;
        }

        key_pos += 19;
        size_t key_end = request.find("\r\n", key_pos);
        string client_key = request.substr(key_pos, key_end - key_pos);

        // Формирование ответа
        string response = "HTTP/1.1 101 Switching Protocols\r\n"
                        "Upgrade: websocket\r\n"
                        "Connection: Upgrade\r\n"
                        "Sec-WebSocket-Accept: " + get_websocket_accept(client_key) + "\r\n\r\n";

        send(client_socket, response.c_str(), response.size(), 0);

        // Цикл обработки сообщений
        while(true) {
            char header[2];
            if(recv(client_socket, header, 2, 0) <= 0) break;

            // Парсинг заголовка
            unsigned char opcode = header[0] & 0x0F;
            bool masked = (header[1] & 0x80) != 0;
            size_t len = header[1] & 0x7F;

            // Чтение расширенной длины
            if(len == 126) {
                char ext_len[2];
                recv(client_socket, ext_len, 2, 0);
                len = (ext_len[0] << 8) | ext_len[1];
            } else if(len == 127) {
                cerr << "64-битная длина не поддерживается" << endl;
                break;
            }

            // Чтение маски
            char mask[4];
            if(masked) recv(client_socket, mask, 4, 0);

            // Чтение данных
            vector<char> data(len);
            recv(client_socket, data.data(), len, 0);
            if(masked) unmask_data(data.data(), len, mask);

            // Обработка opcode
            if(opcode == 0x8) { // Закрытие
                cout << "Закрытие соединения" << endl;
                break;
            }

            // Эхо-ответ
            unsigned char response_header[2] = {0x81, static_cast<char>(len)};
            send(client_socket, reinterpret_cast<char*>(response_header), 2, 0);
            send(client_socket, data.data(), len, 0);
        }

        close(client_socket);
    }

    close(server_fd);
    return 0;
}