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

static int r_all(int fd, void *buf, int n) {
    int t = 0; unsigned char *p = buf;
    while (t < n) { int r = read(fd, p + t, n - t); if (r <= 0) return -1; t += r; }
    return 0;
}

int main(int argc, char *argv[]) {
    int port = argc > 1 ? atoi(argv[1]) : 8080;
    setbuf(stderr, NULL);
    int srv = socket(AF_INET, SOCK_STREAM, 0);
    int opt = 1; setsockopt(srv, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    struct sockaddr_in a = {.sin_family = AF_INET, .sin_port = htons(port)};
    bind(srv, (struct sockaddr *)&a, sizeof(a)); listen(srv, 1);
    fprintf(stderr, "C WS Fast Server on :%d\n", port);

    while (1) {
        socklen_t sl = sizeof(a); int fd = accept(srv, (struct sockaddr *)&a, &sl);
        int one = 1; setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
        fprintf(stderr, "[+] %s:%d\n", inet_ntoa(a.sin_addr), ntohs(a.sin_port));

        char buf[4096]; int t = 0;
        while (t < 4095 && !strstr(buf, "\r\n\r\n")) {
            int r = read(fd, buf + t, 4095 - t); if (r <= 0) goto close; t += r; buf[t] = 0;
        }
        static const char *resp =
            "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n"
            "Connection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n";
        write(fd, resp, strlen(resp));

        unsigned char frame[65536]; int count = 0;
        while (1) {
            unsigned char hdr[2];
            if (r_all(fd, hdr, 2) < 0) break;
            int op = hdr[0] & 0x0F, masked = (hdr[1] >> 7) & 1;
            int plen = hdr[1] & 0x7F;
            if (plen == 126) { unsigned char e[2]; r_all(fd, e, 2); plen = (e[0]<<8)|e[1]; }
            else if (plen == 127) { unsigned char e[8]; r_all(fd, e, 8); plen=0; for(int i=0;i<8;i++)plen=(plen<<8)|e[i]; }
            unsigned char mk[4] = {0};
            if (masked) r_all(fd, mk, 4);
            if (plen > 0) r_all(fd, frame, plen);
            if (masked) for (int i = 0; i < plen; i++) frame[i] ^= mk[i % 4];

            if (op == 0x08) { unsigned char cl[] = {0x88, 0x00}; write(fd, cl, 2); break; }
            if (op == 0x09) { unsigned char po[6] = {0x8A, 0x80}; write(fd, po, 2); continue; }

            if (plen > 0) {
                unsigned char h[10]; int hl = 2; h[0] = 0x80 | op;
                if (plen < 126) h[1] = plen;
                else { h[1] = 126; h[2] = (plen>>8)&0xFF; h[3] = plen&0xFF; hl = 4; }
                struct iovec iov[2] = {{h, hl}, {frame, plen}};
                writev(fd, iov, 2); /* single syscall = single TCP segment */
                count++;
            }
        }
        fprintf(stderr, "[-] %d msgs echoed\n", count);
close:  close(fd);
    }
    return 0;
}