module;

#include <asio.hpp>

#include <expected>
#include <format>

export module ttwhy:io;

namespace ttwhy::io
{

  export template<typename Stream>
  auto handle_signals(Stream & stream) -> asio::awaitable<std::error_code>
  {
    auto executor = co_await asio::this_coro::executor;
    auto signals = asio::signal_set{executor, SIGINT, SIGTERM};

    auto [error, signal] = co_await signals.async_wait(asio::as_tuple(asio::use_awaitable));
    if (!error)
    {
      auto message = std::format("Received signal {}, exiting ...\n", signal);
      co_await asio::async_write(stream, asio::buffer(message), asio::use_awaitable);
    }
    co_return error;
  }

  export template<typename InputStream, typename ErrorStream>
  auto read(std::span<char> data, InputStream & input_stream, ErrorStream & error_stream)
      -> asio::awaitable<std::expected<std::size_t, asio::error_code>>
  {
    auto [error, read] = co_await input_stream.async_read_some(asio::buffer(data), asio::as_tuple(asio::use_awaitable));

    if (error)
    {
      auto message = std::format("{}, exiting ...\n", error.message());
      co_await asio::async_write(error_stream, asio::buffer(message), asio::use_awaitable);
      co_return std::unexpected{error};
    }

    co_return read;
  }

  export template<typename OutputStream, typename ErrorStream>
  auto write(std::span<char const> string, OutputStream & output_stream, ErrorStream & error_stream)
      -> asio::awaitable<asio::error_code>
  {
    auto [error, written] =
        co_await asio::async_write(output_stream, asio::buffer(string), asio::as_tuple(asio::use_awaitable));

    if (error)
    {
      auto message = std::format("{}, exiting ...\n", error.message());
      co_await asio::async_write(error_stream, asio::buffer(message), asio::use_awaitable);
    }

    co_return error;
  }

  export template<typename InputStream, typename OutputStream, typename ErrorStream>
  auto echo(InputStream & input_stream, OutputStream & output_stream, ErrorStream & error_stream)
      -> asio::awaitable<asio::error_code>
  {
    auto input = std::array<char, 16>{};

    while (true)
    {
      auto read_result = co_await read(input, input_stream, error_stream);
      if (!read_result.has_value())
      {
        co_return read_result.error();
      }

      auto to_write = std::span{input.data(), read_result.value()};
      auto write_error = co_await write(to_write, output_stream, error_stream);
      if (write_error)
      {
        co_return write_error;
      }
    }
  }

}  // namespace ttwhy::io