const std = @import("std");
const ArgsTuple = std.meta.ArgsTuple;
const Allocator = std.mem.Allocator;
const ArenaAllocator = std.heap.ArenaAllocator;
const Init = std.process.Init;
const Cmd = @This();
const utils = @import("utils.zig");
const Handler = *const fn (Init, []const []const u8) anyerror!void;

fn displayName(comptime T: type) []const u8 {
    const name = @typeName(T);
    if (name[0] == '?') return displayName(@typeInfo(T).optional.child);
    if (std.mem.eql(u8, name, "[]const u8")) return "string";
    if (name.len > 1 and name[0] == 'f' and isDigit(name[1])) return "float";
    if (name.len > 1 and ((name[0] == 'i' and isDigit(name[1])) or (name[0] == 'u' and isDigit(name[1])))) return "integer";
    return name;
}

fn isDigit(c: u8) bool {
    return c >= '0' and c <= '9';
}

key: []const u8,
value: Value,
description: []const u8 = "",

/// Declare a CLI command from a Zig function. Command name, argument names,
/// descriptions, and defaults are derived from `cli.Param()` types on the
/// function's parameters.
pub inline fn cmd(comptime Fn: anytype, comptime description: []const u8) Cmd {
    const FnType = @TypeOf(Fn);
    const fn_params = @typeInfo(FnType).@"fn".params;

    comptime var cli_count: usize = 0;
    inline for (fn_params) |fp| {
        const T = fp.type orelse @compileError("command function parameters must have explicit types");
        if (T == Init or T == Allocator) continue;
        if (!isParamType(T))
            @compileError("unsupported parameter type: " ++ @typeName(T) ++ " (expected cli.Param)");
        cli_count += 1;
    }

    comptime var param_list: [cli_count]Command.Param = undefined;
    comptime var idx: usize = 0;
    inline for (fn_params, 0..) |fp, fi| {
        const T = fp.type.?;
        if (T == Init or T == Allocator) continue;
        param_list[idx] = .{
            .name = T.param_name,
            .description = T.param_description,
            .type_name = displayName(T.BaseType),
            .values = possibleValues(T.BaseType),
            .required = T.param_type == .req,
        };
        _ = fi;
        idx += 1;
    }
    const frozen = param_list;
    const W = wrap(Fn);

    return .{
        .key = utils.funcName(Fn),
        .description = description,
        .value = .{
            .cmd = .{
                .params = &frozen,
                .Handler = W.wrapper,
                .Validate = W.validate,
            },
        },
    };
}

/// list of exclusive sub commands
pub inline fn list(
    /// Arg name used for the group in the CLI.
    comptime key: []const u8,
    /// Help text for the group (groups have no function, so no doc source).
    comptime description: []const u8,
    /// Args under this group.
    comptime cmds: []const Cmd,
) Cmd {
    return .{
        .key = key,
        .description = description,
        .value = .{ .list = cmds },
    };
}

pub const Command = struct {
    params: []const Command.Param = &.{},
    Handler: Handler = undefined,
    Validate: ?*const fn (Init, []const []const u8, ?*Diag) anyerror!void = null,

    pub const Param = struct {
        name: []const u8,
        description: []const u8,
        type_name: []const u8,
        values: []const u8,
        required: bool,
    };
};

pub const Value = union(enum) {
    cmd: Command,
    list: []const Cmd,
};

/// Check if a type is a `cli.Param(...)` by looking for its comptime declarations.
fn isParamType(comptime T: type) bool {
    return comptime blk: {
        break :blk @hasDecl(T, "param_name") and @hasDecl(T, "BaseType") and
            @hasDecl(T, "param_type") and @hasDecl(T, "param_description");
    };
}

/// Extract possible values from an enum or tagged union at comptime,
/// returning a string like "[installed, all, updates]" or "" for types
/// with no discrete values.
fn possibleValues(comptime T: type) []const u8 {
    return comptime blk: {
        const info = @typeInfo(T);
        const fields = switch (info) {
            .@"enum" => |e| e.fields,
            .optional => |o| break :blk possibleValues(o.child),
            .@"union" => |u| {
                if (u.tag_type) |tag| break :blk possibleValues(tag);
                break :blk "";
            },
            else => break :blk "",
        };
        if (fields.len == 0) break :blk "";
        break :blk "[" ++ joinFieldNames(fields) ++ "]";
    };
}

fn joinFieldNames(comptime fields: []const std.builtin.Type.EnumField) []const u8 {
    if (fields.len == 0) return "";
    if (fields.len == 1) return fields[0].name;
    return fields[0].name ++ ", " ++ joinFieldNames(fields[1..]);
}

/// Diagnostic info about the parameter that failed to parse.
pub const Diag = struct {
    param: []const u8 = "",
    value: []const u8 = "",
};

/// Wraps a function in handlers that parse argv slices into the wrapped
/// function's parameter tuple. `validate` parses and discards the result so
/// argument errors can be reported before the command runs.
fn wrap(comptime f: anytype) type {
    const Fn = @TypeOf(f);
    const fn_params = @typeInfo(Fn).@"fn".params;
    return struct {
        fn parseAll(init: Init, raw_args: []const []const u8, diag: ?*Diag) anyerror!ArgsTuple(Fn) {
            var args: ArgsTuple(Fn) = undefined;
            var arena: ArenaAllocator = .init(init.gpa);
            defer arena.deinit();
            const arena_alloc = arena.allocator();

            var cli_index: usize = 0;
            inline for (fn_params, 0..) |fp, i| {
                const T = fp.type.?;
                if (T == Init) {
                    args[i] = init;
                    continue;
                }
                if (T == Allocator) {
                    args[i] = arena_alloc;
                    continue;
                }

                const P = T;
                const raw = if (cli_index < raw_args.len) raw_args[cli_index] else "";
                cli_index += 1;

                if (raw.len == 0 and P.param_type == .default) {
                    args[i] = .{ .value = P.param_type.default };
                } else if (raw.len == 0) {
                    return error.MissingParameter;
                } else {
                    args[i] = .{ .value = parseArg(P.BaseType, arena_alloc, raw) catch |err| {
                        if (diag) |d| {
                            d.param = P.param_name;
                            d.value = raw;
                        }
                        return err;
                    } };
                }
            }
            return args;
        }

        pub fn validate(init: Init, raw_args: []const []const u8, diag: ?*Diag) anyerror!void {
            if (fn_params.len == 0) return;
            _ = try parseAll(init, raw_args, diag);
        }

        pub fn wrapper(init: Init, raw_args: []const []const u8) anyerror!void {
            if (fn_params.len == 0) return @call(.auto, f, .{});
            const args = try parseAll(init, raw_args, null);
            return @call(.auto, f, args);
        }
    };
}

fn parseArg(comptime T: type, arena: Allocator, raw: []const u8) !T {
    switch (@typeInfo(T)) {
        .optional => |o| {
            if (raw.len == 0) return null;
            return try parseArg(o.child, arena, raw);
        },
        .pointer => |p| {
            if (p.size != .slice) @compileError("only slice pointers supported: " ++ @typeName(T));
            if (p.child == u8 and p.is_const) return raw;
            return try parseList(T, arena, raw);
        },
        .int => return std.fmt.parseInt(T, raw, 10),
        .float => return std.fmt.parseFloat(T, raw),
        .bool => {
            var buf: [5]u8 = undefined;
            const lower = std.ascii.lowerString(&buf, raw);
            for ([_][]const u8{ "true", "yes" }) |s|
                if (std.mem.startsWith(u8, s, lower)) return true;
            for ([_][]const u8{ "false", "no" }) |s|
                if (std.mem.startsWith(u8, s, lower)) return false;
            return error.InvalidBoolean;
        },
        .@"enum" => {
            var partial: ?T = null;
            inline for (std.meta.fields(T)) |field| {
                if (std.mem.eql(u8, field.name, raw))
                    return @enumFromInt(field.value);
                if (std.mem.startsWith(u8, field.name, raw)) {
                    if (partial != null) return error.AmbiguousEnumValue;
                    partial = @enumFromInt(field.value);
                }
            }
            return partial orelse error.InvalidEnumTag;
        },
        else => {
            if (comptime @hasDecl(T, "parse")) {
                const parse_fn = @typeInfo(@TypeOf(T.parse)).@"fn";
                if (parse_fn.params.len == 1) {
                    return try T.parse(raw);
                } else if (parse_fn.params.len == 2) {
                    return try T.parse(arena, raw);
                } else {
                    @compileError("parse function must take ([]const u8) or (Allocator, []const u8)");
                }
            }
            @compileError("unsupported argument type: " ++ @typeName(T));
        },
    }
}

fn parseList(comptime T: type, arena: Allocator, raw: []const u8) !T {
    const Child = @typeInfo(T).pointer.child;
    if (raw.len == 0) return &.{};
    var count: usize = 1;
    for (raw) |c| {
        if (c == ',') count += 1;
    }
    const out = try arena.alloc(Child, count);
    var it = std.mem.splitScalar(u8, raw, ',');
    var i: usize = 0;
    while (it.next()) |part| : (i += 1) {
        out[i] = try parseArg(Child, arena, part);
    }
    return out;
}

/// Base type used for function parameters. Declares a named CLI parameter with a type,
/// required/optional status, and description.
pub fn Param(
    comptime name: []const u8,
    comptime T: type,
    comptime param: Type(T),
    comptime description: []const u8,
) type {
    return struct {
        pub const param_name = name;
        pub const BaseType = T;
        pub const param_type: Type(T) = param;
        pub const param_description = description;

        value: T,
    };
}

/// Optional boolean switch used as a function arg. Value is true if present.
pub fn Switch(comptime name: []const u8, comptime description: []const u8) type {
    return Param(name, bool, .{ .swtch = false }, description);
}

/// List of type T. iterator returns null when out of values.
pub fn List(comptime name: []const u8, comptime T: type, comptime description: []const u8, comptime optional: bool) type {
    return struct {
        const Self = @This();

        // default value for this type is an empty list
        pub const param_name = name;
        pub const BaseType = []T;
        pub const param_type: Type(bool) = .{ .lst = optional };
        pub const param_description = description;

        value: []T,

        pub fn iterator(self: *Self) Iterator {
            _ = self;
            return .{};
        }

        const Iterator = struct {};
    };
}

fn Type(comptime T: type) type {
    return union(enum) {
        req,
        swtch: T,
        lst: T, // if true this arg is optional
        default: T,

        const Self = @This();

        pub const required: Self = .{ .req = {} };
        /// if value is not pulled from args, this one will be used
        pub fn optional(default: T) Self {
            return .{ .default = default };
        }
    };
}