test_exe.zig

const std = @import("std");

pub fn main(init: std.process.Init) !void {
    const allocator = init.gpa;
    var it = try init.minimal.args.iterateAllocator(allocator);
    defer it.deinit();

    var args_list = std.array_list.Managed([]const u8).init(allocator);
    defer args_list.deinit();
    while (it.next()) |arg| {
        try args_list.append(arg);
    }
    const args = args_list.items;

    // detect --stderr, --exit N, --abort and collect positional args (excluding program name)
    var pos_list = std.array_list.Managed([]const u8).init(allocator);
    defer pos_list.deinit();
    var use_stderr = false;
    var i: usize = 1;
    var exit_requested: ?u8 = null;
    var abort_requested = false;
    var print_cwd = false;
    var getenv_name: ?[]const u8 = null;
    var interactive = false;
    while (i < args.len) : (i += 1) {
        const s = args[i];
        if (std.mem.eql(u8, s, "--stderr")) {
            use_stderr = true;
            continue;
        }
        if (std.mem.eql(u8, s, "--interactive")) {
            interactive = true;
            continue;
        }
        if (std.mem.eql(u8, s, "--abort")) {
            abort_requested = true;
            continue;
        }
        if (std.mem.eql(u8, s, "--print-cwd")) {
            print_cwd = true;
            continue;
        }
        if (std.mem.eql(u8, s, "--getenv")) {
            if (i + 1 >= args.len) {
                std.debug.print("missing getenv name\n", .{});
                std.process.exit(1);
            }
            getenv_name = args[i + 1];
            i += 1;
            continue;
        }
        if (std.mem.eql(u8, s, "--exit")) {
            // next argument must be the exit code
            if (i + 1 >= args.len) {
                std.debug.print("missing exit code\n", .{});
                std.process.exit(1);
            }
            const code_s = args[i + 1];
            // parse as unsigned 8-bit integer
            const parsed = std.fmt.parseInt(u8, code_s, 10) catch {
                std.debug.print("invalid exit code: {s}\n", .{code_s});
                std.process.exit(1);
            };
            exit_requested = parsed;
            i += 1; // skip the numeric argument
            continue;
        }
        try pos_list.append(s);
    }
    const pos = pos_list.items;

    // single writer buffer used for whichever stream we pick
    var buf: [1024]u8 = undefined;
    var writer = if (use_stderr)
        std.Io.File.stderr().writer(init.io, &buf)
    else
        std.Io.File.stdout().writer(init.io, &buf);
    const io = &writer.interface;

    if (interactive) {
        var stdin_buffer: [1024]u8 = undefined;
        var stdin_file = std.Io.File.stdin();
        var stdin_reader = stdin_file.reader(init.io, &stdin_buffer);
        var stdin = &stdin_reader.interface;

        while (stdin.takeDelimiterExclusive('\n')) |line| {
            const trimmed = std.mem.trimEnd(u8, line, "\r");
            if (std.mem.eql(u8, trimmed, "PING")) {
                try io.writeAll("PONG\n");
            } else if (std.mem.startsWith(u8, trimmed, "ECHO ")) {
                try io.writeAll(trimmed[5..]);
                try io.writeAll("\n");
            } else if (std.mem.eql(u8, trimmed, "EXIT")) {
                break;
            } else {
                try io.print("UNKNOWN: {s}\n", .{trimmed});
            }
            try io.flush();
            // NOTE: this is required in order to prevents the program loop forever
            _ = try stdin.take(1);
        } else |err| switch (err) {
            error.EndOfStream => {
                std.debug.print("END OF STREAM\n", .{});
            },
            else => |e| return e,
        }

        return;
    }

    if (getenv_name) |name| {
        // Use the process environment map directly from init
        if (init.environ_map.get(name)) |v| {
            try io.writeAll(v);
        }
        try io.writeAll("\n");
    } else if (print_cwd) {
        const cwd = try std.process.currentPathAlloc(init.io, allocator);
        defer allocator.free(cwd);
        try io.writeAll(cwd);
        try io.writeAll("\n");
    } else if (pos.len == 0) {
        try io.writeAll("OK\n");
    } else {
        var j: usize = 0;
        while (j < pos.len) : (j += 1) {
            try io.writeAll(pos[j]);
            try io.writeAll("\n");
        }
    }

    try io.flush();

    // Handle special termination flags after producing output.
    if (abort_requested) {
        // abort will terminate the process with SIGABRT on POSIX.
        std.process.abort();
    }

    if (exit_requested) |code| {
        std.process.exit(code);
    }
}