#include <arpa/inet.h>
#include <netinet/in.h>
#include <openssl/sha.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#define PORT 3456

static const char *GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
static const char b64[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
    "abcdefghijklmnopqrstuvwxyz"
    "0123456789+/";

void base64_encode(
  const unsigned char *in,
  size_t len,
  char *out
) {
  size_t i, j = 0;

  for (i = 0; i < len; i += 3) {
    unsigned v = in[i] << 16;

    if (i + 1 < len) v |= in[i + 1] << 8;
    if (i + 2 < len) v |= in[i + 2];

    out[j++] = b64[(v >> 18) & 63];
    out[j++] = b64[(v >> 12) & 63];

    out[j++] = (i + 1 < len) ? b64[(v >> 6) & 63] : '=';
    out[j++] = (i + 2 < len) ? b64[v & 63] : '=';
  }

  out[j] = 0;
}

void websocket_accept_key(const char *client_key, char *output) {
  char buf[256];
  snprintf(buf, sizeof(buf), "%s%s", client_key, GUID);

  unsigned char hash[SHA_DIGEST_LENGTH];
  SHA1((unsigned char *)buf, strlen(buf), hash);
  base64_encode(hash, SHA_DIGEST_LENGTH, output);
}

char *find_websocket_key(char *req) {
    char *p = strstr(req, "Sec-WebSocket-Key:");

    if (!p) return NULL;

    p += strlen("Sec-WebSocket-Key:");

    while (*p == ' ') p++;

    static char key[128];
    int i = 0;
    while (*p && *p != '\r' && *p != '\n')
        key[i++] = *p++;

    key[i] = 0;
    return key;
}

int main(void) {
  int server_fd = socket(AF_INET, SOCK_STREAM, 0);
  int opt = 1;

  setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
  struct sockaddr_in addr = {
      .sin_family = AF_INET,
      .sin_addr.s_addr = INADDR_ANY,
      .sin_port = htons(PORT)
  };

  bind(server_fd, (struct sockaddr *)&addr, sizeof(addr));
  listen(server_fd, 16);
  printf("Listening on %d\n", PORT);

  for (;;) {
    int client = accept(server_fd, NULL, NULL);

    if (client < 0) continue;
    char req[4096];

    int n = recv(client, req, sizeof(req) - 1, 0);
    if (n <= 0) {
      close(client);
      continue;
    }

    req[n] = 0;

    char *key = find_websocket_key(req);

    if (!key) {
      close(client);
      continue;
    }

    char accept_key[128];

    websocket_accept_key(key, accept_key);
    char response[512];

    snprintf(
      response,
      sizeof(response),
      "HTTP/1.1 101 Switching Protocols\r\n"
      "Upgrade: websocket\r\n"
      "Connection: Upgrade\r\n"
      "Sec-WebSocket-Accept: %s\r\n"
      "\r\n",
      accept_key
    );

    send(client, response, strlen(response), 0);
    printf("WebSocket connected\n");

    for (;;) {
      unsigned char hdr[2];

      if (recv(client, hdr, 2, MSG_WAITALL) != 2)
        break;

      unsigned opcode = hdr[0] & 0x0F;
      unsigned len = hdr[1] & 0x7F;
      if (opcode == 0x8) break;
      if (len >= 126) break;

      unsigned char mask[4];
      recv(client, mask, 4, MSG_WAITALL);

      unsigned char payload[126];
      recv(client, payload, len, MSG_WAITALL);

      for (unsigned i = 0; i < len; i++)
        payload[i] ^= mask[i % 4];

      printf("Message: %.*s\n", len, payload);
      
      unsigned char out[128];
      out[0] = 0x81;
      out[1] = len;

      memcpy(out + 2, payload, len);
      send(client, out, len + 2, 0);
    }
    close(client);
  }
}