#pragma once

#include <exception>
#include <variant>

#include <experimental/coroutine>
#include <experimental/net>

namespace ams::async {
using std::experimental::coroutine_handle;

template <typename Object, typename Function, typename Result>
struct Awaiter {
  Object& _object;
  Function _func;
  std::variant<std::monostate, Result, std::error_code> _result;

  Awaiter(Object& object, Function&& func)
      : _object{object}, _func{std::forward<Function>(func)}, _result{} {}

  bool await_ready() { return false; }

  Result await_resume() {
    if (_result.index() == 2) throw std::system_error(std::get<2>(_result));

    return std::get<1>(std::move(_result));
  }

  void await_suspend(coroutine_handle<> coro) {
    _func(_object, [this, coro](std::error_code const& ec, Result r) mutable {
      if (!ec) {
        _result.template emplace<1>(std::move(r));
      } else {
        _result.template emplace<2>(ec);
      }
      coro.resume();
    });
  }
};

template <typename Object, typename Function>
struct Awaiter<Object, Function, void> {
  Object& _object;
  Function _func;
  std::variant<std::monostate, std::error_code> _result;

  Awaiter(Object& object, Function&& func)
      : _object{object}, _func{std::forward<Function>(func)}, _result{} {}

  bool await_ready() { return false; }

  void await_resume() {
    if (_result.index() == 1) throw std::system_error(std::get<1>(_result));
  }

  void await_suspend(coroutine_handle<> coro) {
    _func(_object, [this, coro](std::error_code const& ec) mutable {
      if (ec) {
        this->_result.template emplace<1>(ec);
      }
      coro.resume();
    });
  }
};

template <typename Result, typename Object, typename Function>
auto make_awaiter(Object& object, Function&& function) {
  return Awaiter<Object, Function, Result>{object,
                                           std::forward<Function>(function)};
}

template <typename Acceptor>
auto async_accept(Acceptor& a) {
  return make_awaiter<typename Acceptor::protocol_type::socket>(
      a, [](Acceptor& acceptor, auto&& callback) {
        acceptor.async_accept(std::forward<decltype(callback)>(callback));
      });
}

template <typename Socket, typename Endpoint>
auto async_connect(Socket& socket, Endpoint ep) {
  return make_awaiter<void>(
      socket, [ep = std::move(ep)](Socket & socket, auto&& callback) {
        socket.async_connect(ep, std::forward<decltype(callback)>(callback));
      });
}

template <typename Socket, typename MutableBuffer>
auto async_read(Socket& socket, MutableBuffer buffer) {
  return make_awaiter<const unsigned long>(
      socket, [buffer = std::move(buffer)](Socket & socket, auto&& callback) {
        std::experimental::net::async_read(
            socket, buffer, std::forward<decltype(callback)>(callback));
      });
}

template <typename Socket, typename ConstantBuffer>
auto async_write(Socket& socket, ConstantBuffer buffer) {
  return make_awaiter<const unsigned long>(
      socket, [buffer = std::move(buffer)](Socket & socket, auto&& callback) {
        std::experimental::net::async_write(
            socket, buffer, std::forward<decltype(callback)>(callback));
      });
}
}  // namespace ams::async