fetch_ml/cli/src/commands/queue/parse.zig

const std = @import("std");

/// Parse job template from command line arguments
pub const JobTemplate = struct {
    job_names: std.ArrayList([]const u8),
    commit_id_override: ?[]const u8,
    priority: u8,
    snapshot_id: ?[]const u8,
    snapshot_sha256: ?[]const u8,
    args_override: ?[]const u8,
    note_override: ?[]const u8,
    cpu: u8,
    memory: u8,
    gpu: u8,
    gpu_memory: ?[]const u8,
    dry_run: bool,
    validate: bool,
    explain: bool,
    json: bool,
    force: bool,
    runner_args_start: ?usize,

    pub fn init(allocator: std.mem.Allocator) JobTemplate {
        return .{
            .job_names = std.ArrayList([]const u8).init(allocator),
            .commit_id_override = null,
            .priority = 5,
            .snapshot_id = null,
            .snapshot_sha256 = null,
            .args_override = null,
            .note_override = null,
            .cpu = 2,
            .memory = 8,
            .gpu = 0,
            .gpu_memory = null,
            .dry_run = false,
            .validate = false,
            .explain = false,
            .json = false,
            .force = false,
            .runner_args_start = null,
        };
    }

    pub fn deinit(self: *JobTemplate, allocator: std.mem.Allocator) void {
        self.job_names.deinit(allocator);
    }
};

/// Parse command arguments into a job template
pub fn parseArgs(allocator: std.mem.Allocator, args: []const []const u8) !JobTemplate {
    var template = JobTemplate.init(allocator);
    errdefer template.deinit(allocator);

    var i: usize = 0;
    while (i < args.len) : (i += 1) {
        const arg = args[i];

        if (std.mem.eql(u8, arg, "--")) {
            template.runner_args_start = i + 1;
            break;
        } else if (std.mem.eql(u8, arg, "--commit-id")) {
            if (i + 1 < args.len) {
                template.commit_id_override = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--priority")) {
            if (i + 1 < args.len) {
                template.priority = std.fmt.parseInt(u8, args[i + 1], 10) catch 5;
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--snapshot")) {
            if (i + 1 < args.len) {
                template.snapshot_id = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--snapshot-sha256")) {
            if (i + 1 < args.len) {
                template.snapshot_sha256 = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--args")) {
            if (i + 1 < args.len) {
                template.args_override = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--note")) {
            if (i + 1 < args.len) {
                template.note_override = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--cpu")) {
            if (i + 1 < args.len) {
                template.cpu = std.fmt.parseInt(u8, args[i + 1], 10) catch 2;
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--memory")) {
            if (i + 1 < args.len) {
                template.memory = std.fmt.parseInt(u8, args[i + 1], 10) catch 8;
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--gpu")) {
            if (i + 1 < args.len) {
                template.gpu = std.fmt.parseInt(u8, args[i + 1], 10) catch 0;
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--gpu-memory")) {
            if (i + 1 < args.len) {
                template.gpu_memory = args[i + 1];
                i += 1;
            }
        } else if (std.mem.eql(u8, arg, "--dry-run")) {
            template.dry_run = true;
        } else if (std.mem.eql(u8, arg, "--validate")) {
            template.validate = true;
        } else if (std.mem.eql(u8, arg, "--explain")) {
            template.explain = true;
        } else if (std.mem.eql(u8, arg, "--json")) {
            template.json = true;
        } else if (std.mem.eql(u8, arg, "--force")) {
            template.force = true;
        } else if (!std.mem.startsWith(u8, arg, "-")) {
            // Positional argument - job name
            try template.job_names.append(arg);
        }
    }

    return template;
}

/// Get runner args from the parsed template
pub fn getRunnerArgs(self: JobTemplate, all_args: []const []const u8) []const []const u8 {
    if (self.runner_args_start) |start| {
        if (start < all_args.len) {
            return all_args[start..];
        }
    }
    return &[_][]const u8{};
}

/// Resolve commit ID from prefix or full hash
pub fn resolveCommitId(allocator: std.mem.Allocator, base_path: []const u8, input: []const u8) ![]u8 {
    if (input.len < 7 or input.len > 40) return error.InvalidArgs;
    for (input) |c| {
        if (!std.ascii.isHex(c)) return error.InvalidArgs;
    }

    if (input.len == 40) {
        return allocator.dupe(u8, input);
    }

    var dir = if (std.fs.path.isAbsolute(base_path))
        try std.fs.openDirAbsolute(base_path, .{ .iterate = true })
    else
        try std.fs.cwd().openDir(base_path, .{ .iterate = true });
    defer dir.close();

    var it = dir.iterate();
    var found: ?[]u8 = null;
    errdefer if (found) |s| allocator.free(s);

    while (try it.next()) |entry| {
        if (entry.kind != .directory) continue;
        const name = entry.name;
        if (name.len != 40) continue;
        if (!std.mem.startsWith(u8, name, input)) continue;
        for (name) |c| {
            if (!std.ascii.isHex(c)) break;
        } else {
            if (found != null) return error.InvalidArgs;
            found = try allocator.dupe(u8, name);
        }
    }

    if (found) |s| return s;
    return error.FileNotFound;
}