const std = @import("std");

fn generateWebSocketKey(allocator: std.mem.Allocator) ![]u8 {
    var random_bytes: [16]u8 = undefined;
    std.crypto.random.bytes(&random_bytes);

    const base64 = std.base64.standard.Encoder;
    const result = try allocator.alloc(u8, base64.calcSize(random_bytes.len));
    _ = base64.encode(result, &random_bytes);
    return result;
}

pub fn handshake(
    allocator: std.mem.Allocator,
    stream: std.net.Stream,
    host: []const u8,
    url: []const u8,
    api_key: []const u8,
) !void {
    const key = try generateWebSocketKey(allocator);
    defer allocator.free(key);

    const request = try std.fmt.allocPrint(
        allocator,
        "GET {s} HTTP/1.1\r\n" ++
            "Host: {s}\r\n" ++
            "Upgrade: websocket\r\n" ++
            "Connection: Upgrade\r\n" ++
            "Sec-WebSocket-Key: {s}\r\n" ++
            "Sec-WebSocket-Version: 13\r\n" ++
            "X-API-Key: {s}\r\n" ++
            "\r\n",
        .{ url, host, key, api_key },
    );
    defer allocator.free(request);

    _ = try stream.write(request);

    var response_buf: [4096]u8 = undefined;
    var bytes_read: usize = 0;
    var header_complete = false;

    while (!header_complete and bytes_read < response_buf.len - 1) {
        const chunk_bytes = try stream.read(response_buf[bytes_read..]);
        if (chunk_bytes == 0) break;
        bytes_read += chunk_bytes;

        if (std.mem.indexOf(u8, response_buf[0..bytes_read], "\r\n\r\n") != null) {
            header_complete = true;
        }
    }

    const response = response_buf[0..bytes_read];

    if (std.mem.indexOf(u8, response, "101 Switching Protocols") == null) {
        if (std.mem.indexOf(u8, response, "404 Not Found") != null) {
            std.debug.print("\nWebSocket Connection Failed\n", .{});
            std.debug.print("-------------------------------------------------------\n\n", .{});
            std.debug.print("The WebSocket endpoint '/ws' was not found on the server.\n\n", .{});
            std.debug.print("This usually means:\n", .{});
            std.debug.print("\t* API server is not running\n", .{});
            std.debug.print("\t* Incorrect server address in config\n", .{});
            std.debug.print("\t* Different service running on that port\n\n", .{});
            std.debug.print("To diagnose:\n", .{});
            std.debug.print("\t* Verify server address: Check ~/.ml/config.toml\n", .{});
            std.debug.print("\t* Test connectivity: curl http://<server>:<port>/health\n", .{});
            std.debug.print("\t* Contact your server administrator if the issue persists\n\n", .{});
            return error.EndpointNotFound;
        } else if (std.mem.indexOf(u8, response, "401 Unauthorized") != null) {
            std.debug.print("\nAuthentication Failed\n", .{});
            std.debug.print("-------------------------------------------------------\n\n", .{});
            std.debug.print("Invalid or missing API key.\n\n", .{});
            std.debug.print("To fix:\n", .{});
            std.debug.print("\t* Verify API key in ~/.ml/config.toml matches server configuration\n", .{});
            std.debug.print("\t* Request a new API key from your administrator if needed\n\n", .{});
            return error.AuthenticationFailed;
        } else if (std.mem.indexOf(u8, response, "403 Forbidden") != null) {
            std.debug.print("\nAccess Denied\n", .{});
            std.debug.print("-------------------------------------------------------\n\n", .{});
            std.debug.print("Your API key doesn't have permission for this operation.\n\n", .{});
            std.debug.print("To fix:\n", .{});
            std.debug.print("\t* Contact your administrator to grant necessary permissions\n\n", .{});
            return error.PermissionDenied;
        } else if (std.mem.indexOf(u8, response, "503 Service Unavailable") != null) {
            std.debug.print("\nServer Unavailable\n", .{});
            std.debug.print("-------------------------------------------------------\n\n", .{});
            std.debug.print("The server is temporarily unavailable.\n\n", .{});
            std.debug.print("This could be due to:\n", .{});
            std.debug.print("\t* Server maintenance\n", .{});
            std.debug.print("\t* High load\n", .{});
            std.debug.print("\t* Server restart\n\n", .{});
            std.debug.print("To resolve:\n", .{});
            std.debug.print("\t* Wait a moment and try again\n", .{});
            std.debug.print("\t* Contact administrator if the issue persists\n\n", .{});
            return error.ServerUnavailable;
        } else {
            std.debug.print("\nWebSocket Handshake Failed\n", .{});
            std.debug.print("-------------------------------------------------------\n\n", .{});
            std.debug.print("Expected HTTP 101 Switching Protocols, but received:\n", .{});

            const newline_pos = std.mem.indexOf(u8, response, "\r\n") orelse response.len;
            const status_line = response[0..newline_pos];
            std.debug.print("  {s}\n\n", .{status_line});

            std.debug.print("To diagnose:\n", .{});
            std.debug.print("\t* Verify server address in ~/.ml/config.toml\n", .{});
            std.debug.print("\t* Check network connectivity to the server\n", .{});
            std.debug.print("\t* Contact your administrator for assistance\n\n", .{});
            return error.HandshakeFailed;
        }
    }

    std.posix.nanosleep(0, 10 * std.time.ns_per_ms);
}
