#ifndef SHARGS_HPP
#define SHARGS_HPP 1

#include <string>
#include <vector>
#include <memory>
#include <optional>
#include <string_view>
#include <sstream>
#include <functional>

class OptionParser
{
private:
    using arg_iterator = std::vector<std::string_view>::iterator;
    using option_return = std::optional<arg_iterator>;
    using unknown_callback = std::function<void(std::string_view)>;

    struct Option {
        const std::string lform;
        const char sform;

        Option(const std::string &l, const char s) : lform(l), sform(s) {}

        /** Given an iterator to the arg after this, fetch arguments and stop
         * when another option switch is encountered.
         * When fetch fails, return the current it.
         * When fetch encounters the end, return nullopt.
         * @return Iterator to next argument
         * */
        virtual option_return parse(arg_iterator it, arg_iterator end) = 0;
    };

    template<typename T>
    struct TypedOption : public Option {
        T &flag;
        TypedOption(const std::string &l, const char s, T &f) :
            Option(l, s), flag(f) {}

         /** Given an iterator to the arg after this, fetch an argument or stop
         * when an option switch is encountered.
         * When fetch is stopped prematurely, return the current it.
         * When fetch encounters the end, return nullopt.
         * @return Iterator to next argument
         * */
        virtual option_return parse(arg_iterator it, arg_iterator end) {
            if(it == end) return std::nullopt;
            std::string_view itval = *it;

            if(itval[0] == '-') return it;

            std::stringstream flagstream;
            flagstream << itval;
            flagstream >> flag;
            return ++it;
        }
    };

    template<typename T>
    struct MultipleOption : public Option {
        std::vector<T> &args;
        MultipleOption(const std::string &l, const char s, std::vector<T> &a) :
            Option(l, s), args(a) {}

         /** Given an iterator to the arg after this, fetch multiple arguments
         * and stop when an option switch is encountered.
         * When fetch is stopped prematurely, return the current it.
         * When fetch encounters the end, return nullopt.
         * @return Iterator to next argument
         * */
        option_return parse(arg_iterator it, arg_iterator end)
        {
            if(it == end) return end;
            while(it != end && (*it)[0] != '-') {
                std::string_view itval = *it;
                std::stringstream argstream;
                T val;

                argstream << itval;
                argstream >> val;
                args.emplace_back(val);
                it++;
            }
            return it;
        }
    };

    template<typename T>
    struct FlagOption : public TypedOption<T> {
        const T set;
        FlagOption(const std::string &l, const char s, T &f, const T set) :
            TypedOption<T>(l, s, f), set(set) {}

         /** Given an iterator to the arg after this, set the flag and return it.
         * @return Iterator to next argument.
         * */
        option_return parse(arg_iterator it, arg_iterator end) 
        {
            this->flag = set;
            return it;
        }
    };

    std::vector<std::unique_ptr<Option>> m_options;
    unknown_callback m_unknown_callback;
public:

    /** Register an option that takes 1 naked argument.
     * @param longform Full name of the option. preceeded by '--'.
     * @param shortform Short name of the option. preceeded by '-'.
     * @param flag Variable to fill when option is encountered.
     * @return self for chaining.
     * */
    template<typename T>
    inline OptionParser &
    add_option_argument(const std::string &longform,
                        char shortform,
                        T &flag)
    {
        m_options.emplace_back(
                std::make_unique<TypedOption<T>>(longform, shortform, flag));
        return *this;
    }

    /** Register an option that sets a variable when encountered.
     * @param longform Full name of the option. preceeded by '--'.
     * @param shortform Short name of the option. preceeded by '-'.
     * @param flag Variable to fill when option is encountered.
     * @param set Value to fill into flag.
     * @return self for chaining.
     * */
    template<typename T>
    inline OptionParser &
    add_option_flag(const std::string &longform, 
                    char shortform, 
                    T &flag,
                    const T set)
    {
        m_options.emplace_back(
                std::make_unique<FlagOption<T>>(longform, shortform, flag, set));
        return *this;
    }

    /** Register an option that takes multiple naked arguments
     * @param longform Full name of the option. preceeded by '--'.
     * @param shortform Short name of the option. preceeded by '-'.
     * @param flag Vector fill with naked argumens.
     * @return self for chaining.
     * */
    template<typename T>
    inline OptionParser &
    add_option_multiple(const std::string &longform,
                        char shortform,
                        std::vector<T> &args)
    {
        m_options.emplace_back(
                std::make_unique<MultipleOption<T>>(longform, shortform, args));
        return *this;
    }

    /** Add a callback function when parser encounters a flagged argument without
     * a known option.
     * @param func Callback to call.
     * */
    inline void
    on_unknown(unknown_callback func)
    {
        m_unknown_callback = func;
    }

    /** Loop through argv, isolating any naked args and parsing those with an
     * option switch ('-' or '--').
     * If the args list is empty, it will return an empty vector.
     * @param argc Number of arguments passed to the program (provided by main)
     * @param argv NULL-terminated array of string pointers. (provided by main)
     * @return Vector of all arguments which were not eaten by options.
     * */
    std::vector<std::string> parse(int argc, char **argv)
    {
        /* Convert argc, argv into vector of string_view. */
        if(argc <= 1) return std::vector<std::string>();
        std::vector<std::string_view> args(argc - 1);
        for(size_t i = 1; i < argc; i++) {
            args[i - 1] = std::string_view(argv[i]);
        }

        /* Fetch iterator for looping through passed values */
        arg_iterator argit = args.begin();

        /* We want to notify the program if the user input any naked values or
         * passed arguments we don't know. */
        std::vector<std::string> naked;
        std::vector<std::string_view> unknown;

        while(argit != args.end()) {
            std::vector<std::vector<std::unique_ptr<Option>>::iterator> to_parse;
            std::string_view arg = *argit;
            if(arg[0] != '-') {
                naked.emplace_back(arg);
                argit++;
                continue;
            }

            auto optionsit = m_options.begin();
            if(arg[1] == '-') {
                /* Look through longform names. */
                arg.remove_prefix(2);
                for(; optionsit != m_options.end(); optionsit++) {
                    if((*optionsit)->lform.empty()) continue;
                    if((*optionsit)->lform == arg) {
                        to_parse.emplace_back(optionsit);
                        break;
                    }
                }
                if(optionsit == m_options.end()) {
                    unknown.emplace_back(*argit);
                }
            }else{
                /* Look through shortform names.
                 * The user should be able to pass as many shortform args
                 * in one switch as they want. */
                arg.remove_prefix(1);
                for(size_t c = 0; c < arg.size(); c++) {

                    /* If the argument is unknown, add it to the unknown list. */
                    bool valid = false;
                    /* Shortform names have arbitrary flags, so we search
                     * through each character. */
                    for(optionsit = m_options.begin(); 
                        optionsit != m_options.end();
                        optionsit++) {
                        /* Shortform name matches, add opt to to_parse. */
                        if(arg[c] == (*optionsit)->sform) {
                            valid = true;
                            to_parse.emplace_back(optionsit);
                            break;
                        }
                    }

                    if(!valid) {
                        unknown.emplace_back((*argit).substr(c + 1, 1));
                    }
                }
            }

            /* Parse all options in to_parse. */
            argit = argit + 1;
            for(auto it = to_parse.begin(); it != to_parse.end(); it++) {
                auto parse_ret = (*(*it))->parse(argit, args.end());
                /* Parsing encountered args.end()*/
                if(parse_ret == std::nullopt) {
                    return naked;
                }
                argit = parse_ret.value();
            }
        }

        /* Go through all unknown args */
        if(m_unknown_callback) {
            for(auto it : unknown) {
                m_unknown_callback(it);
            }
        }

        return naked;
    }
};

#endif