fetch_ml/cli/src/utils/io.zig

const std = @import("std");

fn writeAllFd(fd: std.posix.fd_t, data: []const u8) std.Io.Writer.Error!void {
    var off: usize = 0;
    while (off < data.len) {
        const n = std.posix.write(fd, data[off..]) catch return error.WriteFailed;
        if (n == 0) return error.WriteFailed;
        off += n;
    }
}

fn drainStdout(w: *std.Io.Writer, data: []const []const u8, splat: usize) std.Io.Writer.Error!usize {
    _ = w;
    if (data.len == 0) return 0;
    var written: usize = 0;
    for (data) |chunk| {
        try writeAllFd(std.posix.STDOUT_FILENO, chunk);
        written += chunk.len;
    }
    if (splat > 0) {
        const last = data[data.len - 1];
        var i: usize = 0;
        while (i < splat) : (i += 1) {
            try writeAllFd(std.posix.STDOUT_FILENO, last);
            written += last.len;
        }
    }
    return written;
}

fn drainStderr(w: *std.Io.Writer, data: []const []const u8, splat: usize) std.Io.Writer.Error!usize {
    _ = w;
    if (data.len == 0) return 0;
    var written: usize = 0;
    for (data) |chunk| {
        try writeAllFd(std.posix.STDERR_FILENO, chunk);
        written += chunk.len;
    }
    if (splat > 0) {
        const last = data[data.len - 1];
        var i: usize = 0;
        while (i < splat) : (i += 1) {
            try writeAllFd(std.posix.STDERR_FILENO, last);
            written += last.len;
        }
    }
    return written;
}

const stdout_vtable = std.Io.Writer.VTable{ .drain = drainStdout };
const stderr_vtable = std.Io.Writer.VTable{ .drain = drainStderr };

pub fn stdoutWriter() std.Io.Writer {
    return .{ .vtable = &stdout_vtable, .buffer = &[_]u8{}, .end = 0 };
}

pub fn stderrWriter() std.Io.Writer {
    return .{ .vtable = &stderr_vtable, .buffer = &[_]u8{}, .end = 0 };
}

/// Write a JSON value to stdout
pub fn stdoutWriteJson(value: std.json.Value) !void {
    var buf = std.ArrayList(u8).empty;
    defer buf.deinit(std.heap.page_allocator);
    try writeJSONValue(buf.writer(std.heap.page_allocator), value);
    var stdout_file = std.fs.File{ .handle = std.posix.STDOUT_FILENO };
    try stdout_file.writeAll(buf.items);
    try stdout_file.writeAll("\n");
}

fn writeJSONValue(writer: anytype, v: std.json.Value) !void {
    switch (v) {
        .null => try writer.writeAll("null"),
        .bool => |b| try writer.print("{}", .{b}),
        .integer => |i| try writer.print("{d}", .{i}),
        .float => |f| try writer.print("{d}", .{f}),
        .string => |s| try writeJSONString(writer, s),
        .array => |arr| {
            try writer.writeAll("[");
            for (arr.items, 0..) |item, idx| {
                if (idx > 0) try writer.writeAll(",");
                try writeJSONValue(writer, item);
            }
            try writer.writeAll("]");
        },
        .object => |obj| {
            try writer.writeAll("{");
            var first = true;
            var it = obj.iterator();
            while (it.next()) |entry| {
                if (!first) try writer.writeAll(",");
                first = false;
                try writer.print("\"{s}\":", .{entry.key_ptr.*});
                try writeJSONValue(writer, entry.value_ptr.*);
            }
            try writer.writeAll("}");
        },
        .number_string => |s| try writer.print("{s}", .{s}),
    }
}

fn writeJSONString(writer: anytype, s: []const u8) !void {
    try writer.writeAll("\"");
    for (s) |c| {
        switch (c) {
            '"' => try writer.writeAll("\\\""),
            '\\' => try writer.writeAll("\\\\"),
            '\n' => try writer.writeAll("\\n"),
            '\r' => try writer.writeAll("\\r"),
            '\t' => try writer.writeAll("\\t"),
            else => {
                if (c < 0x20) {
                    var buf: [6]u8 = undefined;
                    buf[0] = '\\';
                    buf[1] = 'u';
                    buf[2] = '0';
                    buf[3] = '0';
                    buf[4] = hexDigit(@intCast((c >> 4) & 0x0F));
                    buf[5] = hexDigit(@intCast(c & 0x0F));
                    try writer.writeAll(&buf);
                } else {
                    try writer.writeAll(&[_]u8{c});
                }
            },
        }
    }
    try writer.writeAll("\"");
}

fn hexDigit(v: u8) u8 {
    return if (v < 10) ('0' + v) else ('a' + (v - 10));
}