/**
 * @file ArgumentParser.cpp
 * @author apio (cloudapio.eu)
 * @brief Command-line argument parser.
 *
 * @copyright Copyright (c) 2023, the Luna authors.
 *
 */

#include <luna/StringBuilder.h>
#include <luna/TypeTraits.h>
#include <os/ArgumentParser.h>
#include <os/File.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/utsname.h>

namespace os
{
    ArgumentParser::ArgumentParser() : m_add_short_help_flag(true), m_add_long_help_flag(true)
    {
    }

    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;
        if (short_flag == 'v') m_add_short_version_flag = false;
        if (long_flag == "version"_sv) m_add_long_version_flag = false;

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

    Result<void> ArgumentParser::add_value_argument(StringView& out, char short_flag, StringView long_flag,
                                                    StringView help)
    {
        ValueArgument 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;
        if (short_flag == 'v') m_add_short_version_flag = false;
        if (long_flag == "version"_sv) m_add_long_version_flag = false;

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

    Result<void> ArgumentParser::set_vector_argument(Vector<StringView>& out, StringView name, StringView fallback,
                                                     bool allow_no_more_flags)
    {
        if (m_vector_argument) return err(EINVAL);
        m_vector_argument = &out;
        m_allow_no_more_flags_after_vector_argument_start = allow_no_more_flags;

        PositionalArgument arg { nullptr, name, false, fallback };

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

    Result<void> ArgumentParser::set_vector_argument(Vector<StringView>& out, StringView name, bool required,
                                                     bool allow_no_more_flags)
    {
        if (m_vector_argument) return err(EINVAL);
        m_vector_argument = &out;
        m_allow_no_more_flags_after_vector_argument_start = allow_no_more_flags;

        PositionalArgument arg { nullptr, name, required, {} };

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

    // Change this every year!
    constexpr auto copyright_text = "Copyright (C) 2023, the Luna authors.";
    constexpr auto license_text = "Licensed under the BSD-2 license <https://opensource.org/license/bsd-2-clause/>";

    void ArgumentParser::add_program_info(ProgramInfo info)
    {
        m_program_info = info;
        m_add_short_version_flag = m_add_long_version_flag = true;
    }

    void ArgumentParser::add_system_program_info(StringView name)
    {
        static utsname info;
        check(uname(&info) == 0);

        ProgramInfo system_info {
            .name = name,
            .version = info.release,
            .copyright = copyright_text,
            .license = license_text,
            .authors = {},
            .package = "Luna system",
        };

        add_program_info(system_info);
    }

    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<bool> ArgumentParser::parse(int argc, char* const* argv)
    {
        StringView program_name = argv[0];

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

        bool is_still_parsing_flags = true;
        bool is_parsing_argument_vector = false;

        Vector<PositionalArgument> positional_args = TRY(m_positional_args.shallow_copy());

        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));
                        return false;
                    }

                    if (m_add_long_version_flag && flag == "version"_sv)
                    {
                        version();
                        return false;
                    }

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

                    if (flag.contains('='))
                    {
                        auto v = TRY(flag.split_view('='));
                        StringView actual_flag = v[0];
                        StringView data = v[1];
                        for (const auto& current : m_value_args)
                        {
                            if (current.long_flag == actual_flag)
                            {
                                *current.out = data;
                                found = true;
                                break;
                            }
                        }
                    }

                    if (found) continue;

                    os::eprintln("%s: unrecognized option '%s'", program_name.chars(), arg.chars());
                    short_usage(program_name);
                    return false;
                }
                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));
                            return false;
                        }

                        if (m_add_short_version_flag && c == 'v')
                        {
                            version();
                            return false;
                        }

                        // 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;

                        os::eprintln("%s: invalid option -- '%c'", program_name.chars(), c);
                        short_usage(program_name);
                        return false;
                    }

                    continue;
                }
            }

            if (is_parsing_argument_vector)
            {
                TRY(m_vector_argument->try_append(arg));
                continue;
            }

            Option<PositionalArgument> current = positional_args.try_dequeue();
            if (!current.has_value()) continue;

            if (!current->out)
            {
                is_parsing_argument_vector = true;
                m_vector_argument->clear();
                TRY(m_vector_argument->try_append(arg));
                if (m_allow_no_more_flags_after_vector_argument_start) is_still_parsing_flags = false;
                continue;
            }

            *current->out = arg;
        }

        if (is_parsing_value_argument)
        {
            os::eprintln("%s: option '-%c' requires an argument", program_name.chars(),
                         current_value_argument->short_flag);
            short_usage(program_name);
            return false;
        }

        if (is_parsing_argument_vector)
        {
            // Fill the positional arguments after the vector using the vector's last elements.
            usize i = 1;
            usize remaining_args = positional_args.size();
            if (remaining_args < m_vector_argument->size()) i = m_vector_argument->size() - remaining_args;
            for (const auto& arg : positional_args)
            {
                if (i >= m_vector_argument->size())
                {
                    if (arg.required)
                    {
                        os::eprintln("%s: required argument '%s' not provided", program_name.chars(), arg.name.chars());
                        short_usage(program_name);
                        return false;
                    }
                    else { *arg.out = arg.fallback; }
                }
                else { *arg.out = m_vector_argument->remove_at(i); }
            }
            check(i >= m_vector_argument->size());
        }
        else
        {
            // Loop through all remaining positional arguments.
            for (const auto& arg : positional_args)
            {
                if (arg.required)
                {
                    os::eprintln("%s: required argument '%s' not provided", program_name.chars(), arg.name.chars());
                    short_usage(program_name);
                    return false;
                }
                else
                {
                    if (arg.out) *arg.out = arg.fallback;
                    else
                    {
                        m_vector_argument->clear();
                        TRY(m_vector_argument->try_append(arg.fallback));
                    }
                }
            }
        }

        return true;
    }

    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())); }
            if (!arg.out) sb.add("..."_sv);
        }

        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() || m_add_long_help_flag || m_add_long_version_flag ||
            m_add_short_help_flag || m_add_short_version_flag)
        {
            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 = "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());
            }
        }

        if (m_add_short_help_flag && m_add_long_help_flag)
            printf("  -h, --%-20s %s\n", "help", "show this help message and exit");
        else if (m_add_long_help_flag)
            printf("      --%-20s %s\n", "help", "show this help message and exit");
        else
            printf("  -%-25c %s\n", 'h', "show this help message and exit");

        if (m_add_short_version_flag && m_add_long_version_flag)
            printf("  -v, --%-20s %s\n", "version", "show version information and exit");
        else if (m_add_long_version_flag)
            printf("      --%-20s %s\n", "version", "show version information and exit");
        else
            printf("  -%-25c %s\n", 'v', "show version information and exit");

        if (m_should_exit_on_bad_usage) exit(0);
        else
            return {};
    }

    Result<void> ArgumentParser::version()
    {
        if (m_program_info.package.is_empty())
            printf("%s %s\n", m_program_info.name.chars(), m_program_info.version.chars());
        else
            printf("%s (%s) %s\n", m_program_info.name.chars(), m_program_info.package.chars(),
                   m_program_info.version.chars());

        printf("%s\n%s\n", m_program_info.copyright.chars(), m_program_info.license.chars());

        if (!m_program_info.authors.is_empty()) printf("\n%s\n", m_program_info.authors.chars());

        if (m_should_exit_on_bad_usage) exit(0);
        else
            return {};
    }

    Result<void> ArgumentParser::short_usage(StringView program_name)
    {
        if (m_add_short_help_flag || m_add_long_help_flag)
            os::eprintln("Try running '%s %s' for more information.", program_name.chars(),
                         m_add_long_help_flag ? "--help" : "-h");

        if (m_should_exit_on_bad_usage) exit(1);
        else
            return {};
    }
}