const std = @import("std");
const Allocator = std.mem.Allocator;
const ArrayList = std.ArrayList;
const HashMap = std.HashMap;

pub const FlagType = enum {
    Boolean,
    String,
    Number,
};

pub const FlagDef = struct {
    name: []const u8,
    aliases: []const []const u8,
    flag_type: FlagType,
    description: []const u8,
    default_value: ?[]const u8 = null,
};

pub const FlagValue = union(FlagType) {
    Boolean: bool,
    String: []const u8,
    Number: i32,
};

pub const ParsedFlags = struct {
    values: HashMap([]const u8, FlagValue, StringContext, std.hash_map.default_max_load_percentage),
    allocator: Allocator,

    const StringContext = struct {
        pub fn hash(self: @This(), s: []const u8) u64 {
            _ = self;
            return std.hash_map.hashString(s);
        }
        pub fn eql(self: @This(), a: []const u8, b: []const u8) bool {
            _ = self;
            return std.mem.eql(u8, a, b);
        }
    };

    pub fn init(allocator: Allocator) ParsedFlags {
        return ParsedFlags{
            .values = HashMap([]const u8, FlagValue, StringContext, std.hash_map.default_max_load_percentage).init(allocator),
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *ParsedFlags) void {
        self.values.deinit();
    }

    pub fn getBool(self: *const ParsedFlags, name: []const u8) bool {
        if (self.values.get(name)) |value| {
            return switch (value) {
                .Boolean => |b| b,
                else => false,
            };
        }
        return false;
    }

    pub fn getString(self: *const ParsedFlags, name: []const u8) ?[]const u8 {
        if (self.values.get(name)) |value| {
            return switch (value) {
                .String => |s| s,
                else => null,
            };
        }
        return null;
    }

    pub fn getNumber(self: *const ParsedFlags, name: []const u8) ?i32 {
        if (self.values.get(name)) |value| {
            return switch (value) {
                .Number => |n| n,
                else => null,
            };
        }
        return null;
    }

    pub fn setValue(self: *ParsedFlags, name: []const u8, value: FlagValue) !void {
        try self.values.put(name, value);
    }
};

pub const CommandFlags = struct {
    flags: []const FlagDef,

    pub fn parse(self: *const CommandFlags, args: []const []const u8, allocator: Allocator) !ParsedFlags {
        var parsed = ParsedFlags.init(allocator);

        for (args) |arg| {
            if (arg.len == 0 or arg[0] != '/') continue;

            const flag_text = arg[1..];
            if (flag_text.len == 0) continue;

            var flag_name: []const u8 = undefined;
            var flag_value: ?[]const u8 = null;

            if (std.mem.indexOf(u8, flag_text, ":")) |colon_pos| {
                flag_name = flag_text[0..colon_pos];
                flag_value = flag_text[colon_pos + 1 ..];
            } else {
                flag_name = flag_text;
            }

            const flag_def = self.findFlag(flag_name) orelse {
                return error.UnknownFlag;
            };

            const value = switch (flag_def.flag_type) {
                .Boolean => FlagValue{ .Boolean = true },
                .String => blk: {
                    const str_value = flag_value orelse flag_def.default_value orelse {
                        return error.MissingFlagValue;
                    };
                    break :blk FlagValue{ .String = try allocator.dupe(u8, str_value) };
                },
                .Number => blk: {
                    const str_value = flag_value orelse flag_def.default_value orelse {
                        return error.MissingFlagValue;
                    };
                    const num_value = std.fmt.parseInt(i32, str_value, 10) catch {
                        return error.InvalidNumber;
                    };
                    break :blk FlagValue{ .Number = num_value };
                },
            };

            try parsed.setValue(flag_def.name, value);
        }

        return parsed;
    }

    fn findFlag(self: *const CommandFlags, name: []const u8) ?*const FlagDef {
        for (self.flags) |*flag_def| {
            if (std.mem.eql(u8, flag_def.name, name)) {
                return flag_def;
            }

            for (flag_def.aliases) |alias| {
                if (alias.len > 0 and alias[0] == '/') {
                    if (std.mem.eql(u8, alias[1..], name)) {
                        return flag_def;
                    }
                } else if (std.mem.eql(u8, alias, name)) {
                    return flag_def;
                }
            }
        }
        return null;
    }

    pub fn getHelp(self: *const CommandFlags, command_name: []const u8, allocator: Allocator) ![]const u8 {
        var help = ArrayList(u8).init(allocator);
        defer help.deinit();

        try help.appendSlice(command_name);
        try help.appendSlice(" - Available flags:\n\n");

        for (self.flags) |flag_def| {
            try help.appendSlice("  /");
            try help.appendSlice(flag_def.name);

            if (flag_def.aliases.len > 0) {
                try help.appendSlice(" (");
                for (flag_def.aliases, 0..) |alias, i| {
                    if (i > 0) try help.appendSlice(", ");
                    try help.appendSlice(alias);
                }
                try help.appendSlice(")");
            }

            switch (flag_def.flag_type) {
                .String => try help.appendSlice(":value"),
                .Number => try help.appendSlice(":number"),
                .Boolean => {},
            }

            try help.appendSlice(" - ");
            try help.appendSlice(flag_def.description);
            try help.appendSlice("\n");
        }

        return help.toOwnedSlice();
    }
};

pub const FlagParseError = error{
    UnknownFlag,
    MissingFlagValue,
    InvalidNumber,
    OutOfMemory,
};