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

namespace os
{
    ArgumentParser::ArgumentParser(bool add_help) : m_add_short_help_flag(add_help), m_add_long_help_flag(add_help)
    {
    }

    void ArgumentParser::add_description(StringView description)
    {
        m_description = description;
    }

    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, StringView help)
    {
        SwitchArgument arg { &out, short_flag, long_flag, help };

        if (short_flag == 'h') m_add_short_help_flag = false;
        if (long_flag == "help"_sv) m_add_long_help_flag = false;

        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, StringView help)
    {
        ValueArgument arg { &out, short_flag, long_flag, value_required, {}, help };

        if (short_flag == 'h') m_add_short_help_flag = false;
        if (long_flag == "help"_sv) m_add_long_help_flag = false;

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

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

        if (short_flag == 'h') m_add_short_help_flag = false;
        if (long_flag == "help"_sv) m_add_long_help_flag = false;

        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;

                    if (m_add_long_help_flag && flag == "help"_sv) { TRY(usage(program_name)); }

                    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());
                    short_usage(program_name);
                }
                else if (looks_like_short_flag(arg))
                {
                    StringView flags = &arg[1];

                    for (usize j = 0; j < flags.length(); j++)
                    {
                        char c = flags[j];
                        bool found = false;

                        if (m_add_short_help_flag && c == 'h') { TRY(usage(program_name)); }

                        // Last flag, this could be a value flag
                        if (j + 1 == flags.length())
                        {
                            for (const auto& current : m_value_args)
                            {
                                if (current.short_flag == ' ') continue;

                                if (current.short_flag == c)
                                {
                                    current_value_argument = current;
                                    is_parsing_value_argument = true;
                                    found = true;
                                    break;
                                }
                            }

                            if (found) continue;
                        }

                        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);
                        short_usage(program_name);
                    }

                    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());
                short_usage(program_name);
            }
            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());
                short_usage(program_name);
            }
            else { *arg.out = arg.fallback; }
        }

        return leftovers;
    }

    Result<void> ArgumentParser::usage(StringView program_name)
    {
        StringBuilder sb;
        TRY(sb.format("Usage: %s <options>"_sv, program_name.chars()));

        for (const auto& arg : m_positional_args)
        {
            if (arg.required) { TRY(sb.format(" %s", arg.name.chars())); }
            else { TRY(sb.format(" [%s]", arg.name.chars())); }
        }

        TRY(sb.add('\n'));

        auto usage_line = TRY(sb.string());
        fputs(usage_line.chars(), stdout);

        if (!m_description.is_empty()) { puts(m_description.chars()); }

        if (m_switch_args.size() || m_value_args.size())
        {
            putchar('\n');
            puts("Options:");
        }

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

                printf("  -%-25c %s\n", arg.short_flag, arg.help.chars());
            }
            else if (arg.short_flag == ' ') { printf("      --%-20s %s\n", arg.long_flag.chars(), arg.help.chars()); }
            else { printf("  -%c, --%-20s %s\n", arg.short_flag, arg.long_flag.chars(), arg.help.chars()); }
        }

        for (const auto& arg : m_value_args)
        {
            StringView value_name;
            if (arg.required) value_name = "VALUE"_sv;
            else
                value_name = "[VALUE]"_sv;

            int field_size = 20 - (int)(arg.long_flag.length() + 1);
            if (field_size < 0) field_size = 0;

            if (arg.long_flag.is_empty())
            {
                if (arg.short_flag == ' ') continue;

                printf("  -%c %-20s %s\n", arg.short_flag, value_name.chars(), arg.help.chars());
            }
            else if (arg.short_flag == ' ')
            {
                printf("      --%s %-*s %s\n", arg.long_flag.chars(), field_size, value_name.chars(), arg.help.chars());
            }
            else
            {
                printf("  -%c, --%s %-*s %s\n", arg.short_flag, arg.long_flag.chars(), field_size, value_name.chars(),
                       arg.help.chars());
            }
        }

        exit(0);
    }

    void ArgumentParser::short_usage(StringView program_name)
    {
        if (m_add_short_help_flag || m_add_long_help_flag)
            fprintf(stderr, "Try running '%s %s' for more information.\n", program_name.chars(),
                    m_add_long_help_flag ? "--help" : "-h");

        exit(1);
    }
}