#include "socket.hpp"
#include "awaiter.hpp"

namespace tinyfiber {
namespace tcp {

using namespace make_result;

Result<Socket> Socket::ConnectTo(const std::string& host, uint16_t port) {
  Socket conn;
  // connect
  auto endpoints = conn.resolver_.resolve(host, std::to_string(port));
  asio::error_code error_code;
  Fiber* fiber = GetCurrentFiber();

  asio::async_connect(
      conn.socket_,
      endpoints,
      [fiber, &error_code](
          const asio::error_code& connect_error_code,
          const asio::ip::tcp::endpoint& endpoint) {

        error_code = connect_error_code;
        GetCurrentScheduler()->Add(fiber);
  });

  GetCurrentScheduler()->Suspend();

  if (!error_code) {
    return Result<Socket>::Ok(std::move(conn));
  } else {
    return Result<Socket>::Fail(error_code);
  }
}

Result<Socket> Socket::ConnectToLocal(uint16_t port) {
  return ConnectTo("localhost", port);
}

Result<size_t> Socket::ReadSome(MutableBuffer buffer) {
  asio::error_code error_code;
  size_t bytes_transferred = 0;
  Fiber* fiber = GetCurrentFiber();

  socket_.async_read_some(buffer, [fiber, &error_code, &bytes_transferred](
      const asio::error_code& read_some_error_code,
      size_t bytes_transferred_) {

    error_code = read_some_error_code;
    bytes_transferred = bytes_transferred_;
    GetCurrentScheduler()->Add(fiber);
  });

  GetCurrentScheduler()->Suspend();
  if (!error_code) {  // some bytes read
    return Result<size_t>::Ok(bytes_transferred);
  } else if (error_code == asio::error::eof) {  // eof reached
    return Result<size_t>::Ok(0);
  } else {  // error occurred
    return Result<size_t>::Fail(error_code);
  }
}

Result<size_t> Socket::Read(MutableBuffer buffer) {
  size_t read = 0;
  while (read < buffer.size()) {
    auto bytes_read = ReadSome(asio::buffer(buffer + read, buffer.size() - read));
    if (!bytes_read.IsOk()) {
      return make_result::Fail(bytes_read.Error());
    }
    read += bytes_read.Value();
    if (bytes_read.Value() == 0) {
      break;
    }
  }
  return make_result::Ok(read);
}

Status Socket::Write(ConstBuffer buffer) {
  asio::error_code error_code;
  Fiber* fiber = GetCurrentFiber();

  asio::async_write(socket_, buffer, [fiber, &error_code](
      const asio::error_code& write_error_code,
      size_t bytes_transferred_) {

    error_code = write_error_code;
    GetCurrentScheduler()->Add(fiber);
  });

  GetCurrentScheduler()->Suspend();

  if (!error_code) {
    return Ok();
  } else {
    return Fail(error_code);
  }
}

Status Socket::ShutdownWrite() {
  // Not implemented
  // Use `shutdown` method of ip::tcp::socket
  asio::error_code error_code;
  socket_.shutdown(asio::ip::tcp::socket::shutdown_send, error_code);
  if (!error_code) {
    return Ok();
  } else {
    return Fail(error_code);
  }
}

Socket::Socket(): socket_(GetCurrentScheduler()->GetIOContext()), resolver_(GetCurrentScheduler()->GetIOContext()) {

}

Socket::Socket(asio::ip::tcp::socket&& socket): socket_(std::move(socket)), resolver_(GetCurrentScheduler()->GetIOContext()) {

}

Acceptor::Acceptor(): acceptor_(GetCurrentScheduler()->GetIOContext()) {

}

Status Acceptor::Listen(uint16_t port) {
  asio::ip::tcp::endpoint ep(asio::ip::tcp::v4(), port);
  asio::error_code error_code;

  acceptor_.open(ep.protocol(), error_code);
  if (error_code) {
    return Fail(error_code);
  }

  acceptor_.set_option(asio::ip::tcp::acceptor::reuse_address(true), error_code);
  if (error_code) {
    return Fail(error_code);
  }

  acceptor_.bind(ep, error_code);
  if (error_code) {
    return Fail(error_code);
  }

  acceptor_.listen();
  return Ok();
}

Result<uint16_t> Acceptor::ListenAvailablePort() {
  auto is_listening = Listen(0);
  if (is_listening.IsOk()) {
    return make_result::Ok(GetPort());
  } else {
    return make_result::Fail(is_listening.Error());
  }
}

Result<Socket> Acceptor::Accept() {
  asio::error_code error_code;
  asio::ip::tcp::socket socket(GetCurrentScheduler()->GetIOContext());
  Fiber* current_fiber = GetCurrentFiber();

  acceptor_.async_accept(socket, [current_fiber, &error_code](
      std::error_code accept_error_code) {

    error_code = accept_error_code;
    GetCurrentScheduler()->Add(current_fiber);
  });

  GetCurrentScheduler()->Suspend();
  if (!error_code) {
    return Result<Socket>::Ok(std::move(socket));
  } else {
    return make_result::Fail(error_code);
  }
}

uint16_t Acceptor::GetPort() const {
  return acceptor_.local_endpoint().port();
}

}  // namespace tcp
}  // namespace tinyfiber