#include "args.hh"

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <format>
#include <iostream>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

namespace {

std::string kEmpty;

class OptionImpl : public Args::OptionArgument {
 public:
  enum Type : uint8_t {
    NoArgument,
    RequiredArgument,
    OptionalArgument,
  };

  OptionImpl(Type type, char shortname, std::string longname, std::string arg,
             std::string help)
      : type(type),
        shortname(shortname),
        longname(std::move(longname)),
        arg(std::move(arg)),
        help(std::move(help)) {}

  const Type type;
  const char shortname;
  const std::string longname;
  const std::string arg;
  const std::string help;

  [[nodiscard]] bool is_set() const override { return is_set_; }

  [[nodiscard]] bool has_argument() const override {
    return value_.has_value();
  }

  [[nodiscard]] const std::string& argument() const override {
    if (value_.has_value())
      return value_.value();
    assert(false);
    return kEmpty;
  }

  void clear() {
    is_set_ = false;
    value_.reset();
  }

  void set_argument(std::string value) {
    assert(type == Type::RequiredArgument || type == Type::OptionalArgument);
    is_set_ = true;
    value_ = std::move(value);
  }

  void set_no_argument() {
    assert(type == Type::NoArgument || type == Type::OptionalArgument);
    is_set_ = true;
    value_.reset();
  }

 private:
  bool is_set_{false};
  std::optional<std::string> value_;
};

class ArgsImpl : public Args {
 public:
  explicit ArgsImpl(std::string prgname) : prgname_(std::move(prgname)) {}

  std::shared_ptr<Option> option(char shortname, std::string longname,
                                 std::string help) override {
    auto opt = std::make_shared<OptionImpl>(
        OptionImpl::Type::NoArgument, shortname, std::move(longname),
        /* arg */ std::string(), std::move(help));
    add(opt);
    return opt;
  }

  std::shared_ptr<OptionArgument> option_argument(char shortname,
                                                  std::string longname,
                                                  std::string arg,
                                                  std::string help,
                                                  bool required) override {
    auto opt = std::make_shared<OptionImpl>(
        required ? OptionImpl::Type::RequiredArgument
                 : OptionImpl::Type::OptionalArgument,
        shortname, std::move(longname), std::move(arg), std::move(help));
    add(opt);
    return opt;
  }

  bool run(int argc, char** argv,
           std::vector<std::string_view>* arguments = nullptr) override {
    last_error_.clear();
    for (auto& opt : options_) {
      opt->clear();
    }

    std::string_view prgname;
    if (prgname_.empty()) {
      if (argc > 0)
        prgname = argv[0];
    } else {
      prgname = prgname_;
    }

    for (int a = 1; a < argc; ++a) {
      assert(argv[a]);
      if (argv[a][0] == '-' && argv[a][1] != '\0') {
        if (argv[a][1] == '-') {
          // long option
          size_t eq = 2;
          while (argv[a][eq] != '=' && argv[a][eq] != '\0')
            ++eq;
          size_t end = eq;
          while (argv[a][end] != '\0')
            ++end;

          if (end == 2) {
            // "--", no more options signal
            if (arguments) {
              for (++a; a < argc; ++a)
                arguments->emplace_back(argv[a]);
            }
            break;
          }

          auto name = std::string_view(argv[a] + 2, eq - 2);
          auto it = long_.find(name);
          if (it == long_.end()) {
            last_error_ =
                std::format("{}: unrecognized option '--{}'", prgname, name);
            return false;
          }
          auto& opt = options_[it->second];

          if (eq < end) {
            // long option with argument after equal sign
            switch (opt->type) {
              case OptionImpl::Type::NoArgument:
                last_error_ =
                    std::format("{}: option '--{}' doesn't allow an argument",
                                prgname, name);
                return false;
              case OptionImpl::Type::RequiredArgument:
              case OptionImpl::Type::OptionalArgument:
                opt->set_argument(
                    std::string(argv[a] + eq + 1, end - (eq + 1)));
                break;
            }
          } else {
            switch (opt->type) {
              case OptionImpl::Type::NoArgument:
              case OptionImpl::Type::OptionalArgument:
                opt->set_no_argument();
                break;
              case OptionImpl::Type::RequiredArgument:
                if (++a >= argc) {
                  last_error_ = std::format(
                      "{}: option '--{}' requires an argument", prgname, name);
                  return false;
                }
                opt->set_argument(argv[a]);
                break;
            }
          }
        } else {
          // short options
          char* current = argv[a] + 1;
          for (; *current; ++current) {
            auto it = short_.find(*current);
            if (it == short_.end()) {
              last_error_ =
                  std::format("{}: invalid option -- '{}'", prgname, *current);
              return false;
            }

            auto& opt = options_[it->second];
            switch (opt->type) {
              case OptionImpl::Type::NoArgument:
              case OptionImpl::Type::OptionalArgument:
                opt->set_no_argument();
                break;
              case OptionImpl::Type::RequiredArgument:
                if (++a >= argc) {
                  last_error_ =
                      std::format("{}: option requires an argument -- '{}'",
                                  prgname, *current);
                  return false;
                }
                opt->set_argument(argv[a]);
                break;
            }
          }
        }
      } else {
        if (arguments)
          arguments->emplace_back(argv[a]);
      }
    }
    return true;
  }

  void print_error(std::ostream& out) const override {
    if (last_error_.empty())
      return;

    out << last_error_ << '\n';
  }

  void print_help(std::ostream& out, uint32_t width = 79) const override {
    if (options_.empty())
      return;

    uint32_t indent = 0;
    const uint32_t max_need = width / 2;
    std::vector<uint32_t> option_need;
    for (auto const& opt : options_) {
      uint32_t need;
      if (opt->longname.empty()) {
        need = 4;  // -O
        switch (opt->type) {
          case OptionImpl::Type::NoArgument:
          case OptionImpl::Type::OptionalArgument:
            break;
          case OptionImpl::Type::RequiredArgument:
            need += 1 + (opt->arg.empty() ? 3 : opt->arg.size());
            break;
        }
      } else {
        need = 8 + opt->longname.size();  // -O, --option
        switch (opt->type) {
          case OptionImpl::Type::NoArgument:
            break;
          case OptionImpl::Type::RequiredArgument:
            // =ARG
            need += 1 + (opt->arg.empty() ? 3 : opt->arg.size());
            break;
          case OptionImpl::Type::OptionalArgument:
            // [=ARG]
            need += 3 + (opt->arg.empty() ? 3 : opt->arg.size());
            break;
        }
      }
      need += 2;  // margin

      option_need.emplace_back(need);
      if (need <= max_need) {
        indent = std::max(indent, need);
      }
    }

    print_wrap(out, width, /* indent */ 0,
               "Mandatory arguments to long options"
               " are mandatory for short options too.");
    auto need_it = option_need.begin();
    for (auto const& opt : options_) {
      if (opt->longname.empty()) {
        out << "  -" << opt->shortname;
        switch (opt->type) {
          case OptionImpl::Type::NoArgument:
          case OptionImpl::Type::OptionalArgument:
            break;
          case OptionImpl::Type::RequiredArgument:
            out << " " << (opt->arg.empty() ? "ARG" : opt->arg);
            break;
        }
      } else {
        if (opt->shortname != '\0') {
          out << "  -" << opt->shortname << ", --";
        } else {
          out << "      --";
        }
        out << opt->longname;
        switch (opt->type) {
          case OptionImpl::Type::NoArgument:
            break;
          case OptionImpl::Type::RequiredArgument:
            out << "=" << (opt->arg.empty() ? "ARG" : opt->arg);
            break;
          case OptionImpl::Type::OptionalArgument:
            out << "=[" << (opt->arg.empty() ? "ARG" : opt->arg) << ']';
            break;
        }
      }

      auto need = *need_it++;
      if (need > max_need) {
        out << '\n';
        if (!opt->help.empty()) {
          print_wrap(out, width, 0, opt->help);
        }
      } else {
        if (opt->help.empty()) {
          out << '\n';
        } else {
          out << "  ";  // add margin, already included in need
          while (need++ < indent)
            out << ' ';
          print_wrap(out, width, indent, opt->help);
        }
      }
    }
  }

 private:
  void add(std::shared_ptr<OptionImpl> opt) {
    if (opt->shortname == '\0' && opt->longname.empty()) {
      assert(false);
    } else {
      auto idx = options_.size();
      if (opt->shortname != '\0')
        short_.emplace(opt->shortname, idx);
      if (!opt->longname.empty())
        long_.emplace(opt->longname, idx);
    }
    options_.emplace_back(std::move(opt));
  }

  static inline bool is_whitespace(char c) { return c == ' ' || c == '\t'; }

  static void print_wrap(std::ostream& out, uint32_t width, uint32_t indent,
                         const std::string& str) {
    if (indent + str.size() <= width) {
      out << str << '\n';
      return;
    }
    if (width <= indent || indent + 10 > width) {
      out << '\n';
      out << str << '\n';
      return;
    }
    const std::string indent_str(indent, ' ');
    const uint32_t avail = width - indent;
    size_t offset = 0;
    while (offset + avail < str.size()) {
      uint32_t i = avail;
      while (i > 0 && !is_whitespace(str[offset + i]))
        --i;
      if (i == 0) {
        out << str.substr(offset, avail - 1);
        out << "-\n";
        offset += avail - 1;
      } else {
        out << str.substr(offset, i);
        out << '\n';
        offset += i;
      }
      out << indent_str;
    }
    out << str.substr(offset);
    out << '\n';
  }

  const std::string prgname_;
  std::vector<std::shared_ptr<OptionImpl>> options_;
  std::map<char, size_t> short_;
  std::map<std::string_view, size_t> long_;
  std::string last_error_;
};

}  // namespace

std::shared_ptr<Args::Option> Args::option(std::string longname,
                                           std::string help) {
  return option(/* shortname */ '\0', std::move(longname), std::move(help));
}

std::shared_ptr<Args::OptionArgument> Args::option_argument(
    std::string longname, std::string arg, std::string help, bool required) {
  return option_argument(/* shortname */ '\0', std::move(longname),
                         std::move(arg), std::move(help), required);
}

std::unique_ptr<Args> Args::create(std::string prgname) {
  return std::make_unique<ArgsImpl>(std::move(prgname));
}