const std = @import("std");
const Io = std.Io;
const File = Io.File;
const Init = std.process.Init;
const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const StringArrayHashMap = std.StringArrayHashMapUnmanaged;
const Cmd = @import("Cmd.zig");
const help = @import("help.zig");
const ArgMap = StringArrayHashMap([]const u8);

pub const Param = Cmd.Param;

/// Parse process arguments and execute the targeted function
pub fn run(
    init: Init,
    /// Description of the CLI tool itself, printed by a bare `<bin> -?`
    comptime description: []const u8,
    /// List of commands and subcommands
    comptime cmds: []const Cmd,
) !void {
    const args = try init.minimal.args.toSlice(init.gpa);
    defer init.gpa.free(args);

    if (args.len == 0)
        return std.debug.print("Binary is somehow missing, wtf did you do?\n", .{});

    var tree: ArrayList([]const u8) = .empty;
    defer tree.deinit(init.gpa);
    try tree.append(init.gpa, args[0]);

    var cmd_list: []const Cmd = cmds;
    var group: ?Cmd = null;
    var command: ?Cmd = null;
    var arg_index: usize = 1;

    while (arg_index < args.len) : (arg_index += 1) {
        const string = args[arg_index];
        if (isHelpFlag(string)) {
            printGroupHelp(init.io, description, group, cmd_list) catch {};
            std.process.exit(0);
        }
        const match = findMatch(cmd_list, string) catch |err| switch (err) {
            error.Ambiguous => {
                var buf: [4096]u8 = undefined;
                var w = File.stderr().writer(init.io, &buf);
                help.printError(&w.interface, "ambiguous command '") catch {};
                w.interface.print("{s}'\n\n", .{string}) catch {};
                help.printGroup(&w.interface, description, group, cmd_list) catch {};
                w.interface.flush() catch {};
                std.process.exit(1);
            },
            error.Unknown => {
                var buf: [4096]u8 = undefined;
                var w = File.stderr().writer(init.io, &buf);
                help.printError(&w.interface, "invalid command '") catch {};
                w.interface.print("{s}'\n\n", .{string}) catch {};
                help.printGroup(&w.interface, description, group, cmd_list) catch {};
                w.interface.flush() catch {};
                std.process.exit(1);
            },
        };
        try tree.append(init.gpa, match.key);
        switch (match.value) {
            .cmd => {
                command = match;
                arg_index += 1;
                break;
            },
            .list => {
                group = match;
                cmd_list = match.value.list;
            },
        }
    }

    var params: ArgMap = .empty;
    defer freeArgMap(init.gpa, &params);

    const cmd_params = if (command) |c| c.value.cmd.params else &[_]Cmd.Command.Param{};

    {
        var positional_index: usize = 0;
        var arg_i = arg_index;

        while (arg_i < args.len) : (arg_i += 1) {
            const string = args[arg_i];
            if (isHelpFlag(string)) {
                if (command) |c| {
                    const cmd_name2 = tree.items[tree.items.len - 1];
                    var buf: [4096]u8 = undefined;
                    var w = File.stderr().writer(init.io, &buf);
                    help.printCommand(&w.interface, cmd_name2, c.description, c.value.cmd.params) catch {};
                    w.interface.flush() catch {};
                } else {
                    printGroupHelp(init.io, description, group, cmd_list) catch {};
                }
                std.process.exit(0);
            }
            if (string.len > 0 and string[0] == '-') break;
            if (positional_index >= cmd_params.len) {
                var buf: [4096]u8 = undefined;
                var w = File.stderr().writer(init.io, &buf);
                help.printError(&w.interface, "too many positional arguments\n\n") catch {};
                if (command) |c| {
                    const cmd_name2 = tree.items[tree.items.len - 1];
                    help.printCommand(&w.interface, cmd_name2, c.description, c.value.cmd.params) catch {};
                }
                w.interface.flush() catch {};
                std.process.exit(1);
            }
            try params.put(init.gpa, cmd_params[positional_index].name, string);
            positional_index += 1;
        }
        arg_index = arg_i;
    }

    while (arg_index < args.len) {
        const string = args[arg_index];
        if (isHelpFlag(string)) {
            if (command) |c| {
                const cmd_name2 = tree.items[tree.items.len - 1];
                var buf: [4096]u8 = undefined;
                var w = File.stderr().writer(init.io, &buf);
                help.printCommand(&w.interface, cmd_name2, c.description, c.value.cmd.params) catch {};
                w.interface.flush() catch {};
            } else {
                printGroupHelp(init.io, description, group, cmd_list) catch {};
            }
            std.process.exit(0);
        }

        const diff_index = std.mem.findDiff(u8, "---", string);
        if (diff_index == null or diff_index.? == 0 or diff_index.? > 2) break;
        arg_index += 1;

        var key: []const u8 = string[diff_index.?..];
        var value: []const u8 = "";
        if (std.mem.findScalar(u8, string[diff_index.?..], '=')) |eq| {
            key = string[diff_index.? .. diff_index.? + eq];
            value = string[diff_index.? + eq + 1 ..];
        } else if (arg_index < args.len) {
            const next = args[arg_index];
            if (next.len == 0 or next[0] != '-') {
                value = next;
                arg_index += 1;
            }
        }

        try params.put(init.gpa, key, value);
    }

    if (arg_index < args.len) {
        var buf: [4096]u8 = undefined;
        var w = File.stderr().writer(init.io, &buf);
        help.printError(&w.interface, "unexpected argument '") catch {};
        w.interface.print("{s}'\n\n", .{args[arg_index]}) catch {};
        if (command) |c| {
            const cmd_name2 = tree.items[tree.items.len - 1];
            help.printCommand(&w.interface, cmd_name2, c.description, c.value.cmd.params) catch {};
        }
        w.interface.flush() catch {};
        std.process.exit(1);
    }

    if (command == null) {
        printGroupHelp(init.io, description, group, cmd_list) catch {};
        std.process.exit(1);
    }

    const cmd_name = tree.items[tree.items.len - 1];
    const resolved = try init.gpa.dupe(Cmd.Command.Param, command.?.value.cmd.params);
    defer init.gpa.free(resolved);

    {
        var it = params.iterator();
        while (it.next()) |e| {
            const provided_name = e.key_ptr.*;
            _ = resolveParam(provided_name, resolved) catch |err| switch (err) {
                error.Unknown => {
                    var buf: [4096]u8 = undefined;
                    var w = File.stderr().writer(init.io, &buf);
                    help.printError(&w.interface, "unknown parameter '--") catch {};
                    w.interface.print("{s}'\n\n", .{provided_name}) catch {};
                    help.printCommand(&w.interface, cmd_name, command.?.description, resolved) catch {};
                    w.interface.flush() catch {};
                    std.process.exit(1);
                },
                error.Ambiguous => {
                    var buf: [4096]u8 = undefined;
                    var w = File.stderr().writer(init.io, &buf);
                    help.printError(&w.interface, "ambiguous parameter '--") catch {};
                    w.interface.print("{s}'\n\n", .{provided_name}) catch {};
                    help.printCommand(&w.interface, cmd_name, command.?.description, resolved) catch {};
                    w.interface.flush() catch {};
                    std.process.exit(1);
                },
            };
        }
    }

    for (resolved) |p| {
        if (!p.required) continue;
        if (!hasParamPrefix(&params, p.name)) {
            var buf: [4096]u8 = undefined;
            var w = File.stderr().writer(init.io, &buf);
            help.printError(&w.interface, "missing required parameter '--") catch {};
            w.interface.print("{s}'\n\n", .{p.name}) catch {};
            help.printCommand(&w.interface, cmd_name, command.?.description, resolved) catch {};
            w.interface.flush() catch {};
            std.process.exit(1);
        }
    }

    const raw = try init.gpa.alloc([]const u8, resolved.len);
    defer init.gpa.free(raw);
    for (resolved, 0..) |p, i| {
        raw[i] = findValue(&params, resolved, p.name);
    }

    var diag: Cmd.Diag = .{};
    command.?.value.cmd.Validate.?(init, raw, &diag) catch |err| switch (err) {
        error.InvalidEnumTag => printParseError(init, cmd_name, "invalid value", diag, resolved, command.?) catch {},
        error.AmbiguousEnumValue => printParseError(init, cmd_name, "ambiguous value", diag, resolved, command.?) catch {},
        error.InvalidBoolean => printParseError(init, cmd_name, "invalid boolean", diag, resolved, command.?) catch {},
        error.InvalidCharacter => printParseError(init, cmd_name, "invalid number", diag, resolved, command.?) catch {},
        error.Overflow => printParseError(init, cmd_name, "value out of range", diag, resolved, command.?) catch {},
        else => return err,
    };

    return command.?.value.cmd.Handler(init, raw);
}

fn printParseError(
    init: Init,
    cmd_name: []const u8,
    kind: []const u8,
    diag: Cmd.Diag,
    resolved: []Cmd.Command.Param,
    full_cmd: Cmd,
) !void {
    var buf: [4096]u8 = undefined;
    var w = File.stderr().writer(init.io, &buf);
    help.printError(&w.interface, kind) catch {};
    w.interface.print(" '{s}' for parameter '--{s}'", .{ diag.value, diag.param }) catch {};
    for (resolved) |p| {
        if (std.mem.eql(u8, p.name, diag.param) and p.values.len > 0) {
            w.interface.print(", expected one of {s}", .{p.values}) catch {};
        }
    }
    w.interface.print("\n\n", .{}) catch {};
    help.printCommand(&w.interface, cmd_name, full_cmd.description, resolved) catch {};
    w.interface.flush() catch {};
    std.process.exit(1);
}

fn printGroupHelp(
    io: Io,
    description: []const u8,
    group: ?Cmd,
    cmd_list: []const Cmd,
) !void {
    var buf: [4096]u8 = undefined;
    var w = File.stderr().writer(io, &buf);
    try help.printGroup(&w.interface, description, group, cmd_list);
    try w.interface.flush();
}

fn isHelpFlag(s: []const u8) bool {
    return std.mem.eql(u8, s, "--help") or
        std.mem.eql(u8, s, "-h") or
        std.mem.eql(u8, s, "--?") or
        std.mem.eql(u8, s, "-?");
}

fn hasParamPrefix(map: *const ArgMap, target: []const u8) bool {
    const single = [_]Cmd.Command.Param{.{ .name = target, .description = "", .type_name = "", .values = "", .required = false }};
    var it = map.iterator();
    while (it.next()) |e| {
        _ = resolveParam(e.key_ptr.*, &single) catch continue;
        return true;
    }
    return false;
}

fn findMatch(list: []const Cmd, input: []const u8) !Cmd {
    var matched: ?Cmd = null;
    for (list) |a| {
        if (std.mem.startsWith(u8, a.key, input)) {
            if (matched != null) return error.Ambiguous;
            matched = a;
        }
    }
    return matched orelse error.Unknown;
}

fn freeArgMap(allocator: Allocator, map: *ArgMap) void {
    map.deinit(allocator);
}

fn resolveParam(prefix: []const u8, params: []const Cmd.Command.Param) !Cmd.Command.Param {
    var matched: ?Cmd.Command.Param = null;
    for (params) |p| {
        if (std.mem.startsWith(u8, p.name, prefix)) {
            if (matched != null) return error.Ambiguous;
            matched = p;
        }
    }
    return matched orelse error.Unknown;
}

fn findValue(map: *const ArgMap, params: []const Cmd.Command.Param, canonical: []const u8) []const u8 {
    var it = map.iterator();
    while (it.next()) |e| {
        _ = resolveParam(e.key_ptr.*, params) catch continue;
        if (std.mem.startsWith(u8, canonical, e.key_ptr.*))
            return e.value_ptr.*;
    }
    return "";
}