/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file micro_session.cc
 */

#include "micro_session.h"

#include <tvm/runtime/crt/rpc_common/framing.h>
#include <tvm/runtime/crt/rpc_common/session.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/registry.h>

#include <algorithm>
#include <chrono>
#include <cstdarg>
#include <memory>
#include <sstream>
#include <string>
#include <utility>

#include "../../support/str_escape.h"
#include "../rpc/rpc_channel.h"
#include "../rpc/rpc_endpoint.h"
#include "../rpc/rpc_session.h"
#include "crt_config.h"

namespace tvm {
namespace runtime {
namespace micro_rpc {

class CallbackWriteStream : public WriteStream {
 public:
  explicit CallbackWriteStream(PackedFunc fsend, ::std::chrono::microseconds write_timeout)
      : fsend_{fsend}, write_timeout_{write_timeout} {}

  ssize_t Write(const uint8_t* data, size_t data_size_bytes) override {
    TVMByteArray bytes;
    bytes.data = (const char*)data;
    bytes.size = data_size_bytes;
    if (write_timeout_ == ::std::chrono::microseconds::zero()) {
      fsend_(bytes, nullptr);
    } else {
      fsend_(bytes, write_timeout_.count());
    }

    return static_cast<ssize_t>(data_size_bytes);
  }

  void PacketDone(bool is_valid) override {}

  void SetWriteTimeout(::std::chrono::microseconds timeout) { write_timeout_ = timeout; }

 private:
  PackedFunc fsend_;
  ::std::chrono::microseconds write_timeout_;
};

class MicroTransportChannel : public RPCChannel {
 public:
  enum class State : uint8_t {
    kReset = 0,               // state entered before the transport has been read or written to.
    kSessionTerminated = 1,   // session is terminated, but transport is alive.
    kSessionEstablished = 2,  // session is alive.
  };

  /*!
   * \brief Construct a new MicroTransportChannel.
   * \param fsend A PackedFunc accepting (data_bytes, timeout_usec) and returning the number of
   *  bytes sent. If a timeout_usec elapses before all data is sent, it should return 0.
   * \param frecv A PackedFunc accepting (num_bytes, timeout_usec) and returning a string containing
   *  the received data. Must not return an empty string, except to indicate a timeout.
   * \param session_start_retry_timeout During session initialization, the session start message is
   *  re-sent after this many microseconds elapse without a reply. If 0, the session start message
   *  is sent only once.
   * \param session_start_timeout Session initialization is considered "timed out" if no reply is
   *  received this many microseconds after the session start is sent. If 0, a session start never
   *  times out.
   * \param session_established_timeout Timeout used for the Recv() function. This is used for
   *  messages sent after a session is already established. If 0, Recv() never times out.
   */
  MicroTransportChannel(PackedFunc fsend, PackedFunc frecv,
                        ::std::chrono::microseconds session_start_retry_timeout,
                        ::std::chrono::microseconds session_start_timeout,
                        ::std::chrono::microseconds session_established_timeout)
      : state_{State::kReset},
        session_start_retry_timeout_{session_start_retry_timeout},
        session_start_timeout_{session_start_timeout},
        session_established_timeout_{session_established_timeout},
        write_stream_{fsend, session_start_timeout},
        framer_{&write_stream_},
        receive_buffer_{new uint8_t[TVM_CRT_MAX_PACKET_SIZE_BYTES], TVM_CRT_MAX_PACKET_SIZE_BYTES},
        session_{&framer_, &receive_buffer_, &HandleMessageReceivedCb, this},
        unframer_{session_.Receiver()},
        did_receive_message_{false},
        frecv_{frecv},
        message_buffer_{nullptr} {}

 private:
  static constexpr const size_t kReceiveBufferSizeBytes = 128;

  /*
   * \brief Receive data until either pf() returns true or a timeout occurs.
   *
   * The condition function is called first, so this function may return without performing a read.
   * Following this call, received data is consumed and frecv_ is invoked until the timeout occurs
   * or the condition function passes.
   *
   * \param pf A condition function that returns true when enough data has been received for the
   *  caller to proceed.
   * \param timeout Pointer to number of microseconds to wait before timing out. If nullptr, no
   *  timeout ever occurs in this function, so it may block forever. If 0, a single non-blocking
   *  read is performed, and any data returned is processed.
   * \return true if the condition passed, false if the timeout expired.
   */
  bool ReceiveUntil(TypedPackedFunc<bool(void)> pf, ::std::chrono::microseconds* timeout) {
    if (pf()) {
      return true;
    }

    auto end_time = ::std::chrono::steady_clock::now();
    if (timeout != nullptr) {
      end_time += *timeout;
    }
    for (;;) {
      if (ConsumeReceivedPayload(pf)) {
        return true;
      }

      ::std::string chunk;
      size_t bytes_needed = unframer_.BytesNeeded();
      CHECK_GT(bytes_needed, 0) << "unframer unexpectedly needs no data";
      if (timeout != nullptr) {
        ::std::chrono::microseconds iter_timeout{
            ::std::max(::std::chrono::microseconds{0},
                       ::std::chrono::duration_cast<::std::chrono::microseconds>(
                           end_time - ::std::chrono::steady_clock::now()))};
        chunk = frecv_(bytes_needed, iter_timeout.count()).operator std::string();
      } else {
        chunk = frecv_(bytes_needed, nullptr).operator std::string();
      }
      pending_chunk_ = chunk;
      if (pending_chunk_.size() == 0) {
        // Timeout occurred
        return false;
      }
    }
  }

  static constexpr const int kNumRandRetries = 10;
  static std::atomic<unsigned int> random_seed;

  inline uint8_t GenerateRandomNonce() {
    // NOTE: this is bad concurrent programming but in practice we don't really expect race
    // conditions here, and even if they occur we don't particularly care whether a competing
    // process computes a different random seed. This value is just chosen pseudo-randomly to
    // form an initial distinct session id. Here we just want to protect against bad loads causing
    // confusion.
    unsigned int seed = random_seed.load();
    if (seed == 0) {
      seed = (unsigned int)time(nullptr);
    }
    uint8_t initial_nonce = 0;
    for (int i = 0; i < kNumRandRetries && initial_nonce == 0; ++i) {
      initial_nonce = rand_r(&seed);
    }
    random_seed.store(seed);
    ICHECK_NE(initial_nonce, 0) << "rand() does not seem to be producing random values";
    return initial_nonce;
  }

  bool StartSessionInternal() {
    using ::std::chrono::duration_cast;
    using ::std::chrono::microseconds;
    using ::std::chrono::steady_clock;

    steady_clock::time_point start_time = steady_clock::now();
    ICHECK_EQ(kTvmErrorNoError, session_.Initialize(GenerateRandomNonce()));
    ICHECK_EQ(kTvmErrorNoError, session_.StartSession());

    if (session_start_timeout_ == microseconds::zero() &&
        session_start_retry_timeout_ == microseconds::zero()) {
      ICHECK(ReceiveUntil([this]() -> bool { return session_.IsEstablished(); }, nullptr))
          << "ReceiveUntil indicated timeout expired, but no timeout set!";
      ICHECK(session_.IsEstablished()) << "Session not established, but should be";
      return true;
    }

    auto session_start_end_time = start_time + session_start_timeout_;
    steady_clock::time_point end_time;
    if (session_start_retry_timeout_ != ::std::chrono::microseconds::zero()) {
      end_time = start_time + session_start_retry_timeout_;
    } else {
      end_time = session_start_end_time;
    }

    while (!session_.IsEstablished()) {
      microseconds time_remaining =
          ::std::max(microseconds{0}, duration_cast<microseconds>(end_time - steady_clock::now()));
      if (ReceiveUntil([this]() -> bool { return session_.IsEstablished(); }, &time_remaining)) {
        break;
      }

      if (session_start_timeout_ != microseconds::zero() && end_time >= session_start_end_time) {
        return false;
      }
      end_time += session_start_retry_timeout_;

      ICHECK_EQ(kTvmErrorNoError, session_.Initialize(GenerateRandomNonce()));
      ICHECK_EQ(kTvmErrorNoError, session_.StartSession());
    }

    return true;
  }

 public:
  bool StartSession() {
    ICHECK(state_ == State::kReset)
        << "MicroSession: state_: expected kReset, got " << uint8_t(state_);

    bool to_return = StartSessionInternal();
    if (to_return) {
      write_stream_.SetWriteTimeout(session_established_timeout_);
    }

    return to_return;
  }

  size_t Send(const void* data, size_t size) override {
    const uint8_t* data_bytes = static_cast<const uint8_t*>(data);
    tvm_crt_error_t err = session_.SendMessage(MessageType::kNormal, data_bytes, size);
    ICHECK(err == kTvmErrorNoError) << "SendMessage returned " << err;

    return size;
  }

  size_t Recv(void* data, size_t size) override {
    size_t num_bytes_recv = 0;
    while (num_bytes_recv < size) {
      if (message_buffer_ != nullptr) {
        num_bytes_recv += message_buffer_->Read(static_cast<uint8_t*>(data), size);
        if (message_buffer_->ReadAvailable() == 0) {
          message_buffer_ = nullptr;
          session_.ClearReceiveBuffer();
        }
        if (num_bytes_recv == size) {
          ICHECK(message_buffer_ == nullptr || message_buffer_->ReadAvailable() > 0);
          return num_bytes_recv;
        }
      }

      did_receive_message_ = false;
      if (session_established_timeout_ == ::std::chrono::microseconds::zero()) {
        ICHECK(ReceiveUntil([this]() -> bool { return did_receive_message_; }, nullptr))
            << "ReceiveUntil timeout expired, but no timeout configured!";
      } else {
        if (!ReceiveUntil([this]() -> bool { return did_receive_message_; },
                          &session_established_timeout_)) {
          std::stringstream ss;
          ss << "MicroSessionTimeoutError: failed to read reply message after timeout "
             << session_established_timeout_.count() / 1e6 << "s";

          throw std::runtime_error(ss.str());
        }
      }
    }

    return num_bytes_recv;
  }

  FrameBuffer* GetReceivedMessage() {
    if (did_receive_message_) {
      did_receive_message_ = false;
      return message_buffer_;
    }

    return nullptr;
  }

 private:
  /*!
   * \brief Consume the entire received payload, unless the pf condition is met halfway through.
   *
   * This function expects pending_chunk_ to contain a chunk of unprocessed packet data. It
   * repeatedly writes the chunk to the Unframer until either a) pf() returns True or b) no more
   * data remains to be written.
   *
   * \param pf A PackedFunc which returns true when ReceiveUntil should return.
   * \returns true if pf() returned true during processing; false otherwise.
   */
  bool ConsumeReceivedPayload(TypedPackedFunc<bool(void)> pf) {
    while (pending_chunk_.size() > 0) {
      size_t bytes_consumed = 0;
      int unframer_error = unframer_.Write((const uint8_t*)pending_chunk_.data(),
                                           pending_chunk_.size(), &bytes_consumed);

      ICHECK(bytes_consumed <= pending_chunk_.size())
          << "consumed " << bytes_consumed << " want <= " << pending_chunk_.size();
      pending_chunk_ = pending_chunk_.substr(bytes_consumed);
      if (unframer_error < 0) {
        LOG(ERROR) << "unframer got error code: " << unframer_error;
      } else {
        if (pf()) {
          return true;
        }
      }
    }

    return false;
  }

  static void HandleMessageReceivedCb(void* context, MessageType message_type, FrameBuffer* buf) {
    static_cast<MicroTransportChannel*>(context)->HandleMessageReceived(message_type, buf);
  }

  void HandleMessageReceived(MessageType message_type, FrameBuffer* buf) {
    size_t message_size_bytes;
    switch (message_type) {
      case MessageType::kStartSessionInit:
        break;

      case MessageType::kStartSessionReply:
        state_ = State::kSessionEstablished;
        break;

      case MessageType::kTerminateSession:
        if (state_ == State::kReset) {
          state_ = State::kSessionTerminated;
        } else if (state_ == State::kSessionTerminated) {
          LOG(FATAL) << "SessionTerminatedError: multiple session-terminated messages received; "
                        "device in reboot loop?";
        } else if (state_ == State::kSessionEstablished) {
          LOG(FATAL) << "SessionTerminatedError: remote device terminated connection";
        }
        break;

      case MessageType::kLog:
        uint8_t message[1024];
        message_size_bytes = buf->ReadAvailable();
        if (message_size_bytes == 0) {
          return;
        } else if (message_size_bytes > sizeof(message) - 1) {
          LOG(ERROR) << "Remote log message is too long to display: " << message_size_bytes
                     << " bytes";
          return;
        }

        ICHECK_EQ(buf->Read(message, sizeof(message) - 1), message_size_bytes);
        message[message_size_bytes] = 0;
        LOG(INFO) << "remote: " << message;
        session_.ClearReceiveBuffer();
        return;

      case MessageType::kNormal:
        did_receive_message_ = true;
        message_buffer_ = buf;
        break;
    }
  }

  State state_;
  ::std::chrono::microseconds session_start_retry_timeout_;
  ::std::chrono::microseconds session_start_timeout_;
  ::std::chrono::microseconds session_established_timeout_;
  CallbackWriteStream write_stream_;
  Framer framer_;
  FrameBuffer receive_buffer_;
  Session session_;
  Unframer unframer_;
  bool did_receive_message_;
  PackedFunc frecv_;
  FrameBuffer* message_buffer_;
  std::string pending_chunk_;
};

std::atomic<unsigned int> MicroTransportChannel::random_seed{0};

TVM_REGISTER_GLOBAL("micro._rpc_connect").set_body([](TVMArgs args, TVMRetValue* rv) {
  MicroTransportChannel* micro_channel =
      new MicroTransportChannel(args[1], args[2], ::std::chrono::microseconds(uint64_t(args[3])),
                                ::std::chrono::microseconds(uint64_t(args[4])),
                                ::std::chrono::microseconds(uint64_t(args[5])));
  if (!micro_channel->StartSession()) {
    std::stringstream ss;
    ss << "MicroSessionTimeoutError: session start handshake failed after " << double(args[4]) / 1e6
       << "s";
    throw std::runtime_error(ss.str());
  }
  std::unique_ptr<RPCChannel> channel(micro_channel);
  auto ep = RPCEndpoint::Create(std::move(channel), args[0], "", args[6]);
  auto sess = CreateClientSession(ep);
  *rv = CreateRPCSessionModule(sess);
});

}  // namespace micro_rpc
}  // namespace runtime
}  // namespace tvm

extern "C" {

void TVMLogf(const char* fmt, ...) {
  va_list args;
  char msg_buf[256];
  va_start(args, fmt);
  vsnprintf(msg_buf, sizeof(msg_buf), fmt, args);
  va_end(args);
  LOG(INFO) << msg_buf;
}

void TVMPlatformAbort(int error_code) { ICHECK(false) << "TVMPlatformAbort: " << error_code; }
}