const std = @import("std");

pub fn sendWebSocketFrame(stream: std.net.Stream, payload: []const u8) !void {
    var frame: [14]u8 = undefined;
    var frame_len: usize = 2;

    // FIN=1, opcode=0x2 (binary), MASK=1
    frame[0] = 0x82 | 0x80;

    // Payload length
    if (payload.len < 126) {
        frame[1] = @as(u8, @intCast(payload.len)) | 0x80;
    } else if (payload.len < 65536) {
        frame[1] = 126 | 0x80;
        frame[2] = @intCast(payload.len >> 8);
        frame[3] = @intCast(payload.len & 0xFF);
        frame_len = 4;
    } else {
        return error.PayloadTooLarge;
    }

    // Generate random mask (4 bytes)
    var mask: [4]u8 = undefined;
    var i: usize = 0;
    while (i < 4) : (i += 1) {
        mask[i] = @as(u8, @intCast(@mod(std.time.timestamp(), 256)));
    }

    @memcpy(frame[frame_len .. frame_len + 4], &mask);
    frame_len += 4;

    _ = try stream.write(frame[0..frame_len]);

    var masked_payload = try std.heap.page_allocator.alloc(u8, payload.len);
    defer std.heap.page_allocator.free(masked_payload);

    for (payload, 0..) |byte, j| {
        masked_payload[j] = byte ^ mask[j % 4];
    }

    _ = try stream.write(masked_payload);
}

pub fn receiveBinaryMessage(stream: std.net.Stream, allocator: std.mem.Allocator) ![]u8 {
    var header: [2]u8 = undefined;
    const header_bytes = try stream.read(&header);
    if (header_bytes < 2) return error.ConnectionClosed;

    // Accept both binary (0x82) and text (0x81) frames
    const opcode = header[0] & 0x0F;
    if (opcode != 0x02 and opcode != 0x01) return error.InvalidFrame;

    const masked = (header[1] & 0x80) != 0;
    var payload_len: usize = header[1] & 0x7F;
    var mask_key: [4]u8 = undefined;

    if (payload_len == 126) {
        var len_bytes: [2]u8 = undefined;
        _ = try stream.read(&len_bytes);
        payload_len = (@as(usize, len_bytes[0]) << 8) | len_bytes[1];
    } else if (payload_len == 127) {
        return error.PayloadTooLarge;
    }

    // Read mask key if frame is masked
    if (masked) {
        _ = try stream.read(&mask_key);
    }

    const payload = try allocator.alloc(u8, payload_len);
    errdefer allocator.free(payload);

    var bytes_read: usize = 0;
    while (bytes_read < payload_len) {
        const n = try stream.read(payload[bytes_read..]);
        if (n == 0) return error.ConnectionClosed;
        bytes_read += n;
    }

    // Unmask payload if needed
    if (masked) {
        for (payload, 0..) |byte, i| {
            payload[i] = byte ^ mask_key[i % 4];
        }
    }

    return payload;
}
