const std = @import("std");

pub const std_options: std.Options = .{
    .log_level = .info,
};

const default_opts = struct {
    /// Whether to start the server or the client side.
    is_server: bool,
    /// IPv4 address to listen (server) on or connect (client) to.
    addr: [:0]const u8,
    /// Message size send between client and server.
    msg_size: usize,
}{ .is_server = true, .addr = "127.0.0.1", .msg_size = 4096 };

/// Parse command line arguments.
fn parse_args() !?@TypeOf(default_opts) {
    var opts = default_opts;
    var args = std.process.args();

    // Skip program name.
    _ = args.skip();

    return while (args.next()) |arg| {
        if (std.mem.eql(u8, arg, "-c")) {
            opts.is_server = false;
        } else if (std.mem.eql(u8, arg, "-a")) {
            opts.addr = args.next() orelse break error.MissingAddressValue;
        } else if (std.mem.eql(u8, arg, "-m")) {
            const bytes = args.next() orelse break error.MissingMessageSizeBytesValue;
            opts.msg_size = try std.fmt.parseInt(usize, bytes, 10);
        } else if (std.mem.eql(u8, arg, "-h")) {
            std.debug.print("Usage: pong [args]\n", .{});
            std.debug.print("  -c          start client, rather than server\n", .{});
            std.debug.print("  -a <addr>   listen/connect IPv4 addr (default: {s})\n", .{default_opts.addr});
            std.debug.print("  -m <bytes>  message size in bytes, client only (default: {})\n", .{default_opts.msg_size});
            std.debug.print("  -h          this help message\n", .{});
            break null;
        } else {
            std.log.err("Unknown argument '{s}'", .{arg});
            break null;
        }
    } else opts;
}

fn run_server(alloc: std.mem.Allocator, addr: std.net.Address) !void {
    // Start a new server.
    var server = try addr.listen(.{ .reuse_address = true });
    defer server.deinit();

    while (true) {
        // Accept a new client.
        std.log.info("server listening on {}", .{addr});
        const client = (try server.accept()).stream;
        defer client.close();

        // Allocate message buffer from client provided message size.
        const buf = blk: {
            // (1) Read message size from client (format: "<bytes>\n").
            var msg_buf: [16]u8 = undefined;
            const msg = try client.reader().readUntilDelimiter(&msg_buf, '\n');
            const msg_size = try std.fmt.parseInt(usize, msg, 10);

            // (2) Allocate actual message buffer.
            std.log.info("allocate message buffer with {} bytes", .{msg_size});
            break :blk try alloc.alloc(u8, msg_size);
        };
        defer alloc.free(buf);

        // Function to receive and loop back message buffer.
        const loop_message = struct {
            fn call(c: std.net.Stream, b: []u8) !void {
                _ = try c.readAll(b);
                try c.writeAll(b);
            }
        }.call;

        // Send message buffer in a loop, on error wait for a new client.
        while (true) {
            loop_message(client, buf) catch break;
        }
    }
}

fn run_client(alloc: std.mem.Allocator, addr: std.net.Address, msg_size: usize) !void {
    // Connect to a server.
    std.log.info("client connecting to {}", .{addr});
    var con = try std.net.tcpConnectToAddress(addr);
    defer con.close();

    // Send message size and allocate message buffer.
    const buf = blk: {
        // (1) Send message size to server (format: "<bytes>\n").
        var msg_buf: [16]u8 = undefined;
        const msg = try std.fmt.bufPrint(&msg_buf, "{}\n", .{msg_size});
        try con.writeAll(msg);

        // (2) Allocate actual message buffer.
        std.log.info("allocate message buffer with {} bytes", .{msg_size});
        break :blk try alloc.alloc(u8, msg_size);
    };
    defer alloc.free(buf);

    // Send message buffer in a loop, and periodically report throughput.
    while (true) {
        var count: usize = 0;
        const start = try std.time.Instant.now();
        var delta: u64 = 0;

        while (delta < std.time.ns_per_s) : (delta = (try std.time.Instant.now()).since(start)) {
            try con.writeAll(buf);
            _ = try con.readAll(buf);
            count += 1;
        }

        const bytes_per_sec = count * 2 * buf.len * (delta / std.time.ns_per_s);
        std.log.info("{:.2}/sec ping-pong messages", .{std.fmt.fmtIntSizeBin(bytes_per_sec)});
    }
}

pub fn main() !void {
    // Command line options.
    const opts = try parse_args() orelse return;
    const addr = try std.net.Address.parseIp4(opts.addr, 8080);

    // Allocator.
    var gpa = std.heap.GeneralPurposeAllocator(.{}).init;
    defer std.debug.assert(gpa.deinit() == .ok);
    const alloc = gpa.allocator();

    // Run server or client.
    try if (opts.is_server)
        run_server(alloc, addr)
    else
        run_client(alloc, addr, opts.msg_size);
}