module;

#include <cassert>
#include <expat.h>

export module routemon:xml;

import std;
import :util;

// XML parsing module.
//
// Makes heavy use of C++20 coroutines. Provides combinators to build
// XML (data format) parsers with. To keep overhead low (and allow for
// HALO/CoroElide), many non-polymorphic definitions are marked inline
// (which is not the default in module units). This also has the added
// benefit of allowing HALO across TU boundaries, which is practically
// necessary to reduce unnecessary allocations when building parsers
// using the provided combinators.

// Ensure that Expat is speaking UTF-8
static_assert(std::is_same_v<XML_Char, char>);

namespace routemon::xml {

  struct qname_view {
    std::string_view ns_uri;
    std::string_view local;

    auto operator==(qname_view const& rhs) const -> bool;
  };

  namespace detail {

    constexpr auto qname_sep = '\xFF';
    auto split_name(char const* name) noexcept -> qname_view {
      auto [l, r] = util::split_on(name, qname_sep);
      if (r) return { .ns_uri = l, .local = *r };
      else return { .ns_uri = std::string_view{}, .local = l };
    }

  } // namespace detail

  class attribute_view_iterator {
    std::string_view default_ns_uri_;
    char const* const* attrs_;

    auto advance() -> void {
      attrs_ += 2;
    }

  public:
    using difference_type = std::ptrdiff_t;
    using value_type = std::pair<qname_view, char const*>;

    struct sentinel {
      friend constexpr auto operator==(attribute_view_iterator const& it, sentinel) noexcept -> bool {
        return !*it.attrs_;
      }
    };

    inline explicit attribute_view_iterator(std::string_view default_ns_uri, char const* const* attrs)
      : default_ns_uri_{default_ns_uri}, attrs_{attrs}
    {}

    inline auto operator*() const -> std::pair<qname_view, char const*> {
      if (!*attrs_)
        throw std::runtime_error{"end of attribute list"};
      auto qname = detail::split_name(attrs_[0]);
      if (qname.ns_uri.empty())
        qname.ns_uri = default_ns_uri_;
      return std::make_pair(qname, attrs_[1]);
    }

    // Pre-increment
    inline auto operator++() -> attribute_view_iterator& {
      advance();
      return *this;
    }

    // Post-increment
    inline auto operator++(int) -> attribute_view_iterator {
      auto pre = *this;
      advance();
      return pre;
    }
  };
  static_assert(std::input_iterator<attribute_view_iterator>);

  class attribute_view : std::ranges::view_base {
    std::string_view default_ns_uri_;
    char const* const* attrs_;

  public:
    inline explicit attribute_view(std::string_view default_ns_uri, char const** attrs)
      : default_ns_uri_{default_ns_uri}, attrs_{const_cast<char const* const*>(attrs)}
    {}

    [[nodiscard]] inline auto begin() const -> attribute_view_iterator {
      return attribute_view_iterator{default_ns_uri_, attrs_};
    }

    [[nodiscard]] inline auto end() const -> attribute_view_iterator::sentinel {
      return {};
    }

    inline auto lookup(qname_view want) -> std::optional<util::not_null<util::lazy_zstring_view>> {
      for (auto const& [name, v] : *this) {
        if (name == want) {
          return util::not_null{util::lazy_zstring_view{v}};
        }
      }
      return std::nullopt;
    }
  };
  static_assert(std::ranges::input_range<attribute_view>);

  template<class T> class promise;

  template<class T>
  struct [[clang::coro_await_elidable, clang::coro_return_type]] parser {
    using promise_type = promise<T>;
    using result_type = promise_type::result_type;
    using handle_type = std::coroutine_handle<promise_type>;

  private:
    handle_type h_;

  public:
    explicit parser(handle_type h)
      : h_{h}
    { assert(h); }

    parser(const parser&) = delete;
    parser(parser&& c) noexcept
      : h_{std::exchange(c.h_, nullptr)}
    {}
    auto operator=(const parser&) -> parser& = delete;
    auto operator=(parser&&) -> parser& = delete;

    [[nodiscard]] auto promise() const -> promise_type& {
      return h_.promise();
    }

    ~parser() {
      if (h_) h_.destroy();
    }
  };

  struct start_element_event {
    qname_view name;
    attribute_view attrs;
  };
  struct end_element_event {
    qname_view name;
  };
  struct character_data_event {
    std::string_view data;
  };
  struct processing_instructions_event {
    util::lazy_zstring_view target;
    util::lazy_zstring_view data;
  };
  struct xml_decl_event {
    util::lazy_zstring_view version;
    util::lazy_zstring_view encoding;
    std::optional<bool> standalone;
  };
  struct eof_event {};
  using event = std::variant<start_element_event,
                             end_element_event,
                             character_data_event,
                             processing_instructions_event,
                             xml_decl_event,
                             eof_event>;
  template<class T>
  concept event_type = requires(event ev) { std::get<T>(ev); };

  class executor;
  using executor_ref = util::not_null<executor*>;

  struct current_event_t {
    executor_ref executor;
  };
  auto current_event(executor_ref executor) -> current_event_t {
    return current_event_t{executor};
  }

  class promise_base {
    executor_ref executor_;
    std::coroutine_handle<promise_base> continuation_ = nullptr;

  public:
    // Not having this constructor marked inline messes with coroutine
    // HALO. (Hours 'wasted': many)
    inline explicit promise_base(executor_ref executor)
      : executor_{executor}
    {}

    [[nodiscard]] inline auto executor() const -> executor& {
      return *executor_;
    }

    inline auto base_handle() -> std::coroutine_handle<promise_base> {
      return std::coroutine_handle<promise_base>::from_promise(*this);
    }

    [[nodiscard]] inline auto continuation() const -> std::coroutine_handle<promise_base> {
      return continuation_;
    }
    inline auto set_continuation(std::coroutine_handle<promise_base> c) -> void {
      continuation_ = c;
    }
  };

  struct position {
    std::size_t line;
    std::size_t col;
  };

  class executor {
    XML_Parser p_;
    std::exception_ptr ex_ = nullptr;
    std::coroutine_handle<promise_base> continuation_ = nullptr;
    std::vector<std::optional<std::string>> default_namespace_;
    std::unordered_map<std::string_view, std::vector<std::string>> namespaces_;
    std::optional<event> ev_;
    bool advance_ = true;

    inline auto try_handle_event(event ev) noexcept -> void {
      assert(!ex_);

      try {
        ev_ = std::move(ev);
      } catch (...) {
        ex_ = std::current_exception();
        return;
      }
      advance_ = false;
      if (!continuation_) {
        // Parser returned (all subparsers are done) and has set the continuation to nullptr.
        if (auto s = XML_StopParser(p_, /* resumable */ false); s != XML_STATUS_OK) {
          ex_ = std::make_exception_ptr(std::runtime_error{"unexpected error when stopping XML parser"});
          return;
        }
        ex_ = std::make_exception_ptr(std::runtime_error{"parser did not consume entire XML document"});
        return;
      }
      continuation_.resume();
      if (ex_) {
        // Not sure if it's useful to report this error.
        std::ignore = XML_StopParser(p_, /* resumable */ false);
      }
    }

    static auto handle_start_element(void* ctx, char const* name, char const** attrs) noexcept -> void {
      auto qname = detail::split_name(name);
      static_cast<executor*>(ctx)->try_handle_event(start_element_event{
          .name = qname,
          .attrs = attribute_view{qname.ns_uri, attrs},
        });
    }
    static auto handle_end_element(void* ctx, char const* name) noexcept -> void {
      static_cast<executor*>(ctx)->try_handle_event(end_element_event{
          .name = detail::split_name(name),
        });
    }
    static auto handle_character_data(void* ctx, char const* s, int len) noexcept -> void {
      static_cast<executor*>(ctx)->try_handle_event(character_data_event{
          .data = std::string_view{s, static_cast<std::size_t>(len)},
        });
    }
    static auto handle_processing_instructions(void* ctx, char const* target, char const* data) noexcept -> void {
      static_cast<executor*>(ctx)->try_handle_event(processing_instructions_event{
          .target = util::lazy_zstring_view{target},
          .data = util::lazy_zstring_view{data},
        });
    }
    static auto handle_external_entity_ref(XML_Parser, char const* /* context */, char const* /* base */, char const* /* system_id */, char const* /* public_id */) noexcept -> int {
      return XML_STATUS_ERROR;
    }
    static auto handle_start_namespace_decl(void* ctx, char const* prefix, char const* uri) noexcept -> void {
      if (prefix) {
        static_cast<executor*>(ctx)->namespaces_[std::string_view{prefix}].emplace_back(uri);
      } else {
        static_cast<executor*>(ctx)->default_namespace_.push_back(uri ? std::make_optional<std::string>(uri) : std::nullopt);
      }
    }
    static auto handle_end_namespace_decl(void* ctx, char const* prefix) noexcept -> void {
      if (prefix) {
        static_cast<executor*>(ctx)->namespaces_[std::string_view{prefix}].pop_back();
      } else {
        static_cast<executor*>(ctx)->default_namespace_.pop_back();
      }
    }
    static auto handle_xml_decl(void * ctx, char const* version, char const* encoding, int standalone) noexcept -> void {
      static_cast<executor*>(ctx)->try_handle_event(xml_decl_event{
          .version = util::lazy_zstring_view{version},
          .encoding = util::lazy_zstring_view{encoding},
          .standalone = standalone < 0 ? std::nullopt : std::make_optional(standalone > 0),
        });
    }

    inline auto advance_flag() -> bool {
      return advance_;
    }
    inline auto set_exception(std::exception_ptr ex) -> void {
      ex_ = std::move(ex);
    }
    [[nodiscard]] inline auto take_exception() -> std::exception_ptr {
      return std::exchange(ex_, nullptr);
    }

    template<class T> friend class promise;

  public:
    executor();

    executor(executor const&) = delete;
    executor(executor&&) = delete;
    auto operator=(executor const&) -> executor& = delete;
    auto operator=(executor&&) -> executor& = delete;

    ~executor();

    inline auto set_continuation(std::coroutine_handle<promise_base> c) -> void {
      continuation_ = c;
    }
    inline auto set_advance_flag() -> void {
      if (!ev_ || !std::holds_alternative<eof_event>(ev_.value())) {
        advance_ = true;
      }
    }
    inline auto event() const -> std::optional<event> const& {
      return ev_;
    }
    inline auto resolve_namespace(std::string_view prefix) -> std::optional<std::string_view> {
      if (auto it = namespaces_.find(prefix); it != namespaces_.end() && !it->second.empty())
        return it->second.back();
      return std::nullopt;
    }
    [[nodiscard]] inline auto position() -> position {
      return {
        .line = XML_GetCurrentLineNumber(p_),
        .col = XML_GetCurrentColumnNumber(p_),
      };
    }

    auto start() -> void;
    auto read(std::string_view xml, bool is_final) -> void;
    auto end() -> void;
  };

  template<class T>
  class promise_returnable : public promise_base {
    std::optional<T> returned_value_;

  public:
    using result_type = T;
    using promise_base::promise_base;

    template<class U>
    auto return_value(U&& v) -> void {
      returned_value_.emplace(std::forward<U>(v));
    }
    auto returned_value() -> T&& {
      if (!returned_value_)
        throw std::runtime_error{"XML coroutine did not return"};
      return std::forward<T>(returned_value_.value());
    }
  };

  template<>
  class promise_returnable<void> : public promise_base {
  public:
    using result_type = void;
    using promise_base::promise_base;

    auto return_void() -> void {}
  };

  template<class T>
  class promise : public promise_returnable<T> {
  public:
    // Called with all the coroutine's arguments.
    // Ignoring all but the first argument, which should be the executor.
    template<class... Args>
    explicit promise(executor_ref executor, Args&&...)
      : promise_returnable<T>{executor}
    {}

    auto handle() -> std::coroutine_handle<promise<T>> {
      return {parser<T>::handle_type::from_promise(*this)};
    }

    auto get_return_object() -> parser<T> {
      return parser<T>{handle()};
    }

    auto initial_suspend() {
      return std::suspend_always{};
    }
    auto final_suspend() noexcept {
      struct awaiter {
        std::coroutine_handle<> h_;

        [[nodiscard]] constexpr auto await_ready() const noexcept -> bool { return false; }
        auto await_suspend(std::coroutine_handle<>) -> std::coroutine_handle<> { return h_; }
        constexpr auto await_resume() const noexcept -> void { return; }
      };
      if (this->continuation()) {
        return awaiter{this->continuation()};
      } else {
        this->executor().set_continuation(nullptr);
        return awaiter{std::noop_coroutine()};
      }
    }

    auto unhandled_exception() -> void {
      this->executor().set_exception(std::current_exception());
    }

    auto await_transform(current_event_t const& req) {
      struct awaiter {
        executor_ref executor_;

        [[nodiscard]] constexpr auto await_ready() const noexcept -> bool {
          return !executor_->advance_flag();
        }
        auto await_suspend(std::coroutine_handle<promise<T>> h) -> void {
          executor_->set_continuation(h.promise().base_handle());
        }
        [[nodiscard]] auto await_resume() const -> event {
          assert(executor_->event());
          return executor_->event().value();
        }
      };
      return awaiter{req.executor};
    }

    template<class U>
    auto await_transform(parser<U> const& coro) {
      struct [[clang::coro_await_elidable]] awaiter {
        util::not_null<promise<U>*> next_;

        [[nodiscard]] constexpr auto await_ready() const noexcept -> bool { return false; }
        auto await_suspend(std::coroutine_handle<promise<T>> h) -> std::coroutine_handle<> {
          // Passed coroutine handle will be the same as parser<T>::handle_type::from_promise(*this)
          next_->set_continuation(h.promise().base_handle());
          return next_->handle();
        }
        auto await_resume() -> U {
          // Promise is still valid since coroutine frame is still alive (and suspended):
          // control was transferred back to this coroutine via symmetric transfer in
          // final_suspend(). Assuming that the destructor for coro still needs to run.
          if (auto ex = next_->executor().take_exception()) {
            std::rethrow_exception(ex);
          } else {
            if constexpr (!std::is_void_v<U>) {
              return std::move(next_->returned_value());
            }
          }
        }
      };
      return awaiter{util::not_null{&coro.promise()}};
    }
  };

  // Helpers for handling XML documents. Non-polymorphic functions
  // should be marked inline to allow HALO across TU boundaries.

  template<class T>
  concept unconstrained = true;

  template<class T, template<class U> concept C>
  concept parser_of = requires {
    typename T::result_type;
    requires std::same_as<parser<typename T::result_type>, T>;
      requires C<typename T::result_type>;
  };

  template<parser_of<unconstrained> T>
  using parser_result_t = T::result_type;

  template<class T, class... Args>
  concept parser_invocable = std::invocable<T, Args...> && parser_of<std::invoke_result_t<T, Args...>, unconstrained>;

  template<class Fn, class... Args>
  requires parser_invocable<Fn, Args...>
  using parser_invoke_result_t = parser_result_t<std::invoke_result_t<Fn, Args...>>;

  template<event_type T> auto expect_event(executor_ref e) -> parser<T> {
    auto ev = co_await current_event(e);
    if (!std::holds_alternative<T>(ev)) {
      auto pos = e->position();
      throw std::runtime_error{std::format("at {}:{}: unexpected event type, have {}", pos.line, pos.col, ev.index())};
    }
    e->set_advance_flag();
    co_return std::get<T>(ev);
  }

  inline auto expect_start_element(executor_ref e, qname_view want) -> parser<attribute_view> {
    auto ev = co_await expect_event<start_element_event>(e);
    if (ev.name != want) {
      auto pos = e->position();
      throw std::runtime_error{std::format("at {}:{}: unexpected element started", pos.line, pos.col)};
    }
    co_return ev.attrs;
  }

  inline auto allow_start_element(executor_ref e, qname_view want) -> parser<std::optional<attribute_view>> {
    auto ev = co_await current_event(e);
    if (auto const* pev = std::get_if<start_element_event>(&ev)) {
      if (pev->name == want) {
        e->set_advance_flag();
        co_return pev->attrs;
      }
    }
    co_return std::nullopt;
  }

  inline auto expect_end_element(executor_ref e, qname_view want) -> parser<void> {
    auto ev = co_await expect_event<end_element_event>(e);
    if (ev.name != want) {
      throw std::runtime_error{"unexpected element ended"};
    }
  }

  inline auto ignore_whitespace(executor_ref e) -> parser<void> {
    auto all_whitespace = [](std::string_view s) -> bool {
      for (auto c : s)
        if (c != ' ' && c != '\r' && c != '\n' && c != '\t')
          return false;
      return true;
    };

    while (true) {
      auto ev = co_await current_event(e);
      if (auto const* pev = std::get_if<character_data_event>(&ev); pev && all_whitespace(pev->data)) {
        e->set_advance_flag();
      } else {
        co_return;
      }
    }
  }

  auto expect_element(executor_ref e, qname_view want, parser_invocable<executor_ref, attribute_view> auto p)
    -> parser<parser_invoke_result_t<decltype(p), executor_ref, attribute_view>>
  {
    co_await ignore_whitespace(e);
    auto attrs = co_await expect_start_element(e, want);
    auto&& res = co_await p(e, attrs);
    co_await expect_end_element(e, want);
    co_await ignore_whitespace(e);
    co_return std::forward<parser_invoke_result_t<decltype(p), executor_ref, attribute_view>>(res);
  }

  auto allow_element(executor_ref e, qname_view want, parser_invocable<executor_ref, attribute_view> auto p)
    -> parser<std::optional<parser_invoke_result_t<decltype(p), executor_ref, attribute_view>>>
  {
    co_await ignore_whitespace(e);
    if (auto mattrs = co_await allow_start_element(e, want)) {
      auto&& res = co_await p(e, *mattrs);
      co_await expect_end_element(e, want);
      co_await ignore_whitespace(e);
      co_return std::make_optional(std::forward<parser_invoke_result_t<decltype(p), executor_ref, attribute_view>>(res));
    }
    co_return std::nullopt;
  }

  auto allow_element(executor_ref e, qname_view want, parser_invocable<executor_ref, attribute_view> auto p)
    -> parser<bool>
    requires std::is_void_v<parser_invoke_result_t<decltype(p), executor_ref, attribute_view>>
  {
    co_await ignore_whitespace(e);
    if (auto mattrs = co_await allow_start_element(e, want)) {
      co_await p(e, *mattrs);
      co_await expect_end_element(e, want);
      co_await ignore_whitespace(e);
      co_return true;
    }
    co_return false;
  }

  inline auto ignore_contents(executor_ref e, std::optional<qname_view> muntil = std::nullopt) -> parser<void> {
    std::size_t depth = 0;
    while (true) {
      auto ev = co_await current_event(e);
      if (auto* pev = std::get_if<start_element_event>(&ev)) {
        if (depth == 0 && muntil && pev->name == *muntil) {
          co_return;
        } else {
          depth++;
        }
      } else if (std::holds_alternative<end_element_event>(ev)) {
        if (depth == 0) {
          co_return;
        } else {
          depth--;
        }
      }
      e->set_advance_flag();
    }
  }
  inline auto ignore_element_contents(executor_ref e, attribute_view) -> parser<void> {
    co_await ignore_contents(e);
  }

  inline auto read_string(executor_ref e) -> parser<std::string> {
    std::string s;
    while (true) {
      auto ev = co_await current_event(e);
      if (auto* pev = std::get_if<character_data_event>(&ev)) {
        e->set_advance_flag();
        s += pev->data;
      } else {
        co_return s;
      }
    }
  }
  inline auto read_string_contents(executor_ref e, attribute_view) -> parser<std::string> {
    co_return co_await read_string(e);
  }

  template<auto f>
  auto hohalo() {
    return []<class... Args>(Args&&... args) -> std::invoke_result_t<decltype(f), Args...> {
      co_return co_await f(std::forward<Args>(args)...);
    };
  }

} // namespace routemon::xml