Luna/libos/src/ArgumentParser.cpp

#include <os/ArgumentParser.h>
#include <stdio.h>
#include <stdlib.h>

Result<void> ArgumentParser::add_positional_argument(StringView& out, StringView name, bool required)
{
    PositionalArgument arg { &out, name, required, {} };

    return m_positional_args.try_append(move(arg));
}

Result<void> ArgumentParser::add_positional_argument(StringView& out, StringView name, StringView fallback)
{
    PositionalArgument arg { &out, name, false, fallback };

    return m_positional_args.try_append(move(arg));
}

Result<void> ArgumentParser::add_switch_argument(bool& out, char short_flag, StringView long_flag)
{
    SwitchArgument arg { &out, short_flag, long_flag };

    return m_switch_args.try_append(move(arg));
}

Result<void> ArgumentParser::add_value_argument(StringView& out, char short_flag, StringView long_flag,
                                                bool value_required)
{
    ValueArgument arg { &out, short_flag, long_flag, value_required, {} };

    return m_value_args.try_append(move(arg));
}

Result<void> ArgumentParser::add_value_argument(StringView& out, char short_flag, StringView long_flag,
                                                StringView fallback)
{
    ValueArgument arg { &out, short_flag, long_flag, false, fallback };

    return m_value_args.try_append(move(arg));
}

static bool looks_like_short_flag(StringView arg)
{
    return arg.length() > 1 && arg[0] == '-';
}

static bool looks_like_long_flag(StringView arg)
{
    return arg.length() > 2 && arg[0] == '-' && arg[1] == '-';
}

Result<Vector<StringView>> ArgumentParser::parse(int argc, char* const* argv)
{
    StringView program_name = argv[0];

    Vector<StringView> leftovers;

    Option<ValueArgument> current_value_argument = {};
    bool is_parsing_value_argument = false;

    bool is_still_parsing_flags = true;

    for (int i = 1; i < argc; i++)
    {
        StringView arg = argv[i];

        if (is_parsing_value_argument)
        {
            *current_value_argument->out = arg;
            is_parsing_value_argument = false;
            continue;
        }

        if (is_still_parsing_flags)
        {
            if (arg == "--")
            {
                is_still_parsing_flags = false;
                continue;
            }

            if (looks_like_long_flag(arg))
            {
                StringView flag = &arg[2];

                bool found = false;

                for (const auto& current : m_switch_args)
                {
                    if (current.long_flag == flag)
                    {
                        *current.out = true;
                        found = true;
                        break;
                    }
                }

                for (const auto& current : m_value_args)
                {
                    if (current.long_flag == flag)
                    {
                        current_value_argument = current;
                        is_parsing_value_argument = true;
                        found = true;
                        break;
                    }
                }

                if (found) continue;

                fprintf(stderr, "%s: unrecognized option '%s'\n", program_name.chars(), arg.chars());
                exit(1);
            }
            else if (looks_like_short_flag(arg))
            {
                StringView flags = &arg[1];

                for (char c : flags)
                {
                    bool found = false;

                    // FIXME: Implement value arguments for short flags.

                    for (const auto& current : m_switch_args)
                    {
                        if (current.short_flag == ' ') continue;

                        if (current.short_flag == c)
                        {
                            *current.out = true;
                            found = true;
                            break;
                        }
                    }

                    if (found) continue;

                    fprintf(stderr, "%s: invalid option -- '%c'\n", program_name.chars(), c);
                    exit(1);
                }

                continue;
            }
        }

        Option<PositionalArgument> current = m_positional_args.try_dequeue();
        if (!current.has_value())
        {
            TRY(leftovers.try_append(arg));
            continue;
        }

        *current->out = arg;
    }

    if (is_parsing_value_argument)
    {
        if (current_value_argument->required)
        {
            fprintf(stderr, "%s: option '--%s' requires an argument\n", program_name.chars(),
                    current_value_argument->long_flag.chars());
            exit(1);
        }
        else { *current_value_argument->out = current_value_argument->fallback; }
    }

    // Loop through all remaining positional arguments.
    for (const auto& arg : m_positional_args)
    {
        if (arg.required)
        {
            fprintf(stderr, "%s: required argument '%s' not provided\n", program_name.chars(), arg.name.chars());
            exit(1);
        }
        else { *arg.out = arg.fallback; }
    }

    return leftovers;
}