#include "common.h"

#include <chrono>
#include <array>

#include "a2netcompat.h"
#include "a2io.h"
#include "a2functional.h"
#include "a2time.h"
#include "message.h"
#include "DlRetryEx.h"
#include "DlAbortEx.h"
#include "fmt.h"
#include "util.h"
#include "TimeA2.h"
#include "Peer.h"
#include "Logger.h"
#include "LogFactory.h"
#include "Option.h"
#include "DownloadEngine.h"
#include "SocketCore.h"
#include "uTPCore.h"
#include "RequestGroupMan.h"
#include "ReceiverMSEHandshakeCommand.h"

namespace aria2 {

utp_context* uTPCore::ctx_;
utp_context* uTPCore::ctx6_;
std::shared_ptr<SocketCore> uTPCore::sock_;
std::shared_ptr<SocketCore> uTPCore::sock6_;
Timer uTPCore::start_;

uint64 uTPCore::getMilliseconds(utp_callback_arguments* arg)
{
  auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(
      start_.difference());
  return (uint64)ms.count();
}

uint64 uTPCore::getMicroseconds(utp_callback_arguments* arg)
{
  auto us = std::chrono::duration_cast<std::chrono::microseconds>(
      start_.difference());
  return (uint64)us.count();
}

uint64 uTPCore::getReadBufferSize(utp_callback_arguments* arg)
{
  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
  if (utp == nullptr) {
    return 0;
  }
  return utp->end_ - utp->pos_;
}

uint64 uTPCore::writeLogCallback(utp_callback_arguments* arg)
{
  A2_LOG_DEBUG(fmt("uTP: %s", arg->buf));
  return 0;
}

uint64 uTPCore::sendtoCallback(utp_callback_arguments* arg)
{
  if (arg->address->sa_family == AF_INET) {
    sock_->writeData(arg->buf, arg->len, arg->address, arg->address_len,
                     arg->flags);
  }
  else {
    sock6_->writeData(arg->buf, arg->len, arg->address, arg->address_len,
                      arg->flags);
  }
  return 0;
}

uint64 uTPCore::acceptCallback(utp_callback_arguments* arg)
{
  std::shared_ptr<uTPCore> utp = std::make_shared<uTPCore>(arg->socket);
  utp_set_userdata(arg->socket, utp.get());
  utp->isWritable_ = true;
  Endpoint endpoint = util::getNumericNameInfo(arg->address, arg->address_len);
  DownloadEngine* e = (DownloadEngine*)utp_context_get_userdata(arg->context);
  std::shared_ptr<Peer> peer =
      std::make_shared<Peer>(endpoint.addr, endpoint.port, true);
  cuid_t cuid = e->newCUID();
  e->addCommand(make_unique<ReceiverMSEHandshakeCommand>(cuid, peer, e, utp));
  A2_LOG_DEBUG(fmt("Accepted the connection from %s:%u.",
                   peer->getIPAddress().c_str(), peer->getPort()));
  A2_LOG_DEBUG(
      fmt("Added CUID#%" PRId64 " to receive BitTorrent/MSE handshake.", cuid));
  return 0;
}

uint64 uTPCore::stateChangeCallback(utp_callback_arguments* arg)
{
  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
  if (arg->state == UTP_STATE_CONNECT || arg->state == UTP_STATE_WRITABLE) {
    utp->isWritable_ = true;
    if (utp->writeCommand_ != nullptr && utp->writeCheck_ == true) {
      utp->writeCommand_->setStatusActive();
      utp->writeCommand_->writeEventReceived();
    }
  }
  return 0;
}

uint64 uTPCore::readCallback(utp_callback_arguments* arg)
{
  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
  size_t bufLen = utp->buf_.size() - utp->end_;
  if (arg->len > bufLen) {
    utp->buf_.resize(utp->end_ + arg->len);
  }
  std::memcpy(utp->buf_.data() + utp->end_, arg->buf, arg->len);
  utp->end_ += arg->len;
  if (utp->end_ == utp->buf_.size()) {
    utp->needDrainRead_ = true;
  }
  else {
    utp_read_drained(utp->s_);
  }
  if (utp->readCommand_ != nullptr && utp->readCheck_ == true) {
    utp->readCommand_->setStatusActive();
    utp->readCommand_->readEventReceived();
  }
  return 0;
}

uint64 uTPCore::errorCallback(utp_callback_arguments* arg)
{
  uTPCore* utp = (uTPCore*)utp_get_userdata(arg->socket);
  if (utp != nullptr) {
    utp->errNum_ = arg->error_code + 1;
    if (utp->readCommand_ != nullptr) {
      utp->readCommand_->setStatusActive();
      utp->readCommand_->errorEventReceived();
    }
    if (utp->writeCommand_ != nullptr) {
      utp->writeCommand_->setStatusActive();
      utp->writeCommand_->errorEventReceived();
    }
  }
  return 0;
}

namespace {
void setupContextAndSock(utp_context*& ctx, std::shared_ptr<SocketCore>& sock,
                         int family, uint16_t port, DownloadEngine* e)
{
  sock = std::make_shared<SocketCore>(SOCK_DGRAM);
  sock->bind("", port, family);
  sock->setNonBlockingMode();
  ctx = utp_init(2);
  utp_context_set_userdata(ctx, (void*)e);
  utp_set_callback(ctx, UTP_GET_MILLISECONDS, uTPCore::getMilliseconds);
  utp_set_callback(ctx, UTP_GET_MICROSECONDS, uTPCore::getMicroseconds);
  utp_set_callback(ctx, UTP_GET_READ_BUFFER_SIZE, uTPCore::getReadBufferSize);
  if (A2_LOG_DEBUG_ENABLED) {
    utp_context_set_option(ctx, UTP_LOG_DEBUG, 1);
    utp_context_set_option(ctx, UTP_LOG_NORMAL, 1);
    utp_context_set_option(ctx, UTP_LOG_MTU, 1);
  }
  utp_set_callback(ctx, UTP_LOG, uTPCore::writeLogCallback);
  utp_set_callback(ctx, UTP_SENDTO, uTPCore::sendtoCallback);
  utp_set_callback(ctx, UTP_ON_ACCEPT, uTPCore::acceptCallback);
  utp_set_callback(ctx, UTP_ON_STATE_CHANGE, uTPCore::stateChangeCallback);
  utp_set_callback(ctx, UTP_ON_READ, uTPCore::readCallback);
  utp_set_callback(ctx, UTP_ON_ERROR, uTPCore::errorCallback);

  e->addRoutineCommand(
      make_unique<uTPListenCommand>(e->newCUID(), e, ctx, sock));
}
} // namespace

void uTPCore::initContext(uint16_t port, DownloadEngine* e)
{
  start_ = global::wallclock();
  setupContextAndSock(ctx_, sock_, AF_INET, port, e);
  if (e->getOption()->getAsBool(PREF_DISABLE_IPV6) == false) {
    setupContextAndSock(ctx6_, sock6_, AF_INET6, port, e);
  }
}

void uTPCore::destroyContext()
{
  if (ctx_ != nullptr) {
    utp_destroy(ctx_);
    ctx_ = nullptr;
  }
  if (ctx6_ != nullptr) {
    utp_destroy(ctx6_);
    ctx6_ = nullptr;
  }
}

uTPCore::uTPCore(utp_socket* s)
    : s_(s),
      isWritable_(false),
      needDrainRead_(false),
      pos_(0),
      end_(0),
      errNum_(0),
      readCheck_(false),
      writeCheck_(false),
      readCommand_(nullptr),
      writeCommand_(nullptr),
      SocketCore(SOCK_DGRAM)
{
  buf_.resize(16_k);
}

uTPCore::~uTPCore() { closeConnection(); }

void uTPCore::establishConnection(const std::string& host, uint16_t port,
                                  bool tcpNodelay)
{
  addrinfo* res;
  int rv;
  rv = callGetaddrinfo(&res, host.c_str(), util::uitos(port).c_str(),
                       protocolFamily_, SOCK_DGRAM, 0, 0);
  if (rv != 0) {
    throw DL_ABORT_EX(fmt(EX_RESOLVE_HOSTNAME, host.c_str(), gai_strerror(rv)));
  }
  std::unique_ptr<addrinfo, decltype(&freeaddrinfo)> resDeleter(res,
                                                                freeaddrinfo);
  struct addrinfo* rp;
  for (rp = res; rp; rp = rp->ai_next) {
    s_ = utp_create_socket(rp->ai_family == AF_INET ? ctx_ : ctx6_);
    if (utp_connect(s_, rp->ai_addr, rp->ai_addrlen) != 0) {
      utp_close(s_);
      s_ = nullptr;
      continue;
    }
    utp_set_userdata(s_, (void*)this);
    return;
  }
  throw DL_ABORT_EX(
      fmt("Failed to create uTP connection to %s.", host.c_str()));
}

bool uTPCore::isWritable(time_t timeout) { return isWritable_; }

bool uTPCore::isReadable(time_t timeout) { return pos_ < end_; }

ssize_t uTPCore::writeData(const void* data, size_t len)
{
  a2iovec iov = {(void*)data, len};
  return writeData(&iov, 1);
}

ssize_t uTPCore::writeVector(a2iovec* iov, size_t iovcnt)
{
  if (errNum_ != 0) {
    throw DL_RETRY_EX(
        fmt("Failed to send data, cause: %s.", getSocketError().c_str()));
  }
  ssize_t ret = 0;
  wantRead_ = false;
  wantWrite_ = false;
  size_t len =
      std::accumulate(iov, iov + iovcnt, (size_t)0,
                      [](size_t s, a2iovec& elem) { return s + elem.iov_len; });
  ret = utp_writev(s_, (utp_iovec*)iov, iovcnt);
  if (ret == -1) {
    throw DL_RETRY_EX("Failed to send data over uTP.");
  }
  if ((size_t)ret < len) {
    wantWrite_ = true;
    isWritable_ = false;
  }
  return ret;
}

void uTPCore::readData(void* data, size_t& len)
{
  if (errNum_ != 0) {
    throw DL_RETRY_EX(
        fmt("Failed to read data, cause: %s.", getSocketError().c_str()));
  }
  wantRead_ = false;
  wantWrite_ = false;
  size_t dataLen = end_ - pos_;
  if (dataLen == 0) {
    len = 0;
    wantRead_ = true;
    return;
  }
  if (len > dataLen) {
    len = dataLen;
  }
  std::memcpy(data, buf_.data() + pos_, len);
  pos_ += len;
  if (pos_ > buf_.size() / 2) {
    dataLen = end_ - pos_;
    std::memmove(buf_.data(), buf_.data() + pos_, dataLen);
    pos_ = 0;
    end_ = dataLen;
    if (needDrainRead_ == true) {
      utp_read_drained(s_);
      needDrainRead_ = false;
    }
  }
}

void uTPCore::closeConnection()
{
  if (s_ != nullptr) {
    utp_set_userdata(s_, nullptr);
    utp_close(s_);
  }
  s_ = nullptr;
  readCommand_ = nullptr;
  writeCommand_ = nullptr;
}

std::string uTPCore::getSocketError() const
{
  if (errNum_ == 0) {
    return "NO_ERROR";
  }
  else if (errNum_ >= 1 && errNum_ <= 3) {
    return utp_error_code_names[errNum_ - 1];
  }
  return "UNKNOWN_ERROR";
};

void uTPCore::setRWCheck(Command* command, int is_add, int is_read)
{
  if (is_add == true) {
    if (is_read == true) {
      readCheck_ = true;
      readCommand_ = command;
      if (isReadable(0)) {
        readCommand_->setStatusActive();
        readCommand_->readEventReceived();
      }
    }
    else {
      writeCheck_ = true;
      writeCommand_ = command;
      if (isWritable(0)) {
        writeCommand_->setStatusActive();
        writeCommand_->writeEventReceived();
      }
    }
  }
  else {
    if (is_read == true) {
      readCheck_ = false;
      readCommand_ = nullptr;
    }
    else {
      writeCheck_ = false;
      writeCommand_ = nullptr;
    }
  }
}

uTPListenCommand::uTPListenCommand(cuid_t cuid, DownloadEngine* e,
                                   utp_context* ctx,
                                   std::shared_ptr<SocketCore> sock)
    : e_(e), ctx_(ctx), sock_(sock), Command(cuid)
{
  e_->addSocketForReadCheck(sock_, this);
}

uTPListenCommand::~uTPListenCommand()
{
  if (sock_ != nullptr) {
    e_->deleteSocketForReadCheck(sock_, this);
  }
};

bool uTPListenCommand::execute()
{
  if (e_->isHaltRequested() || e_->getRequestGroupMan()->downloadFinished()) {
    return true;
  }
  std::array<char, 2_k> data;
  try {
    while (1) {
      sockaddr_union sockaddr;
      socklen_t socklen = sizeof(sockaddr);
      ssize_t length =
          sock_->readDataFrom(data.data(), data.size(), sockaddr, socklen);
      if (length == 0) {
        utp_issue_deferred_acks(ctx_);
        break;
      }
      utp_process_udp(ctx_, (byte*)data.data(), (size_t)length, &sockaddr.sa,
                      socklen);
    }
  }
  catch (RecoverableException& e) {
    A2_LOG_INFO_EX("Exception thrown while receiving UDP message.", e);
  }
  utp_check_timeouts(ctx_);
  e_->addCommand(std::unique_ptr<Command>(this));
  return false;
}

} // namespace aria2