#include "mtran/translator_node.hpp"

#include <algorithm>
#include <cctype>
#include <chrono>
#include <functional>
#include <memory>
#include <string>
#include <utility>

namespace mtran
{

TranslatorNode::TranslatorNode(
  const rclcpp::NodeOptions & options,
  std::shared_ptr<TranslationClient> client)
: Node("mtran_bridge_node", options),
  client_(std::move(client)),
  startup_started_(std::chrono::steady_clock::now()),
  request_timeout_ms_(declare_parameter<int>("request_timeout_ms", 30000)),
  startup_timeout_ms_(declare_parameter<int>("startup_timeout_ms", 60000)),
  health_retry_ms_(declare_parameter<int>("health_retry_ms", 500)),
  keep_warm_interval_s_(declare_parameter<int>("keep_warm_interval_s", 43200)),
  max_input_characters_(declare_parameter<int>("max_input_characters", 512))
{
  const auto server_url = declare_parameter<std::string>(
    "server_url", "http://127.0.0.1:8989");
  if (!client_) {
    client_ = std::make_shared<HttpTranslationClient>(server_url, request_timeout_ms_);
  }

  health_retry_ms_ = std::max(1, health_retry_ms_);
  keep_warm_interval_s_ = std::max(1, keep_warm_interval_s_);
  max_input_characters_ = std::max(1, max_input_characters_);

  health_timer_ = create_wall_timer(
    std::chrono::milliseconds(health_retry_ms_),
    std::bind(&TranslatorNode::poll_readiness, this));
}

void TranslatorNode::poll_readiness()
{
  if (ready_) {
    if (!client_->health()) {
      RCLCPP_WARN(get_logger(), "MTranServer health check failed; entering recovery");
      mark_not_ready();
    }
    return;
  }

  if (!ever_ready_ && !startup_timeout_logged_ && startup_timeout_ms_ > 0 &&
    std::chrono::steady_clock::now() - startup_started_ >=
    std::chrono::milliseconds(startup_timeout_ms_))
  {
    startup_timeout_logged_ = true;
    RCLCPP_ERROR(
      get_logger(), "MTranServer did not become ready within %d ms; continuing to retry",
      startup_timeout_ms_);
  }

  if (!client_->health()) {
    return;
  }

  const auto warmup = client_->translate("Hello.");
  if (!warmup.success) {
    return;
  }

  ready_ = true;
  ever_ready_ = true;
  create_service_if_needed();

  if (!keep_warm_timer_) {
    keep_warm_timer_ = create_wall_timer(
      std::chrono::seconds(keep_warm_interval_s_),
      std::bind(&TranslatorNode::keep_warm, this));
  } else {
    keep_warm_timer_->reset();
  }

  RCLCPP_INFO(get_logger(), "MTranServer is healthy and the en-to-zh-Hans model is warm");
}

void TranslatorNode::keep_warm()
{
  if (!ready_) {
    return;
  }

  const auto result = client_->translate("Hello.");
  if (!result.success && result.error == ClientError::kUnavailable) {
    RCLCPP_WARN(get_logger(), "MTranServer keep-warm connection failed; entering recovery");
    mark_not_ready();
  } else if (!result.success) {
    RCLCPP_WARN(get_logger(), "MTranServer keep-warm request failed; keeping service ready");
  }
}

void TranslatorNode::mark_not_ready()
{
  ready_ = false;
  if (keep_warm_timer_) {
    keep_warm_timer_->cancel();
  }
  health_timer_->reset();
}

void TranslatorNode::create_service_if_needed()
{
  if (service_) {
    return;
  }

  service_ = create_service<TranslateService>(
    "/translate_en_to_zh",
    std::bind(
      &TranslatorNode::handle_request, this,
      std::placeholders::_1, std::placeholders::_2));
}

void TranslatorNode::handle_request(
  const TranslateService::Request::SharedPtr request,
  TranslateService::Response::SharedPtr response)
{
  response->success = false;
  response->translation.clear();
  response->error.clear();

  if (is_blank(request->text)) {
    response->error = "INVALID_INPUT";
    return;
  }
  if (utf8_codepoint_count(request->text) >
    static_cast<std::size_t>(max_input_characters_))
  {
    response->error = "INPUT_TOO_LONG";
    return;
  }
  if (!ready_) {
    response->error = "NOT_READY";
    return;
  }

  const auto result = client_->translate(request->text);
  if (result.success) {
    response->success = true;
    response->translation = result.translation;
    return;
  }

  response->error = error_code(result.error);
  if (result.error == ClientError::kUnavailable) {
    mark_not_ready();
  }
}

std::size_t TranslatorNode::utf8_codepoint_count(const std::string & text)
{
  return static_cast<std::size_t>(std::count_if(
           text.begin(), text.end(), [](unsigned char byte) {
             return (byte & 0xC0U) != 0x80U;
           }));
}

bool TranslatorNode::is_blank(const std::string & text)
{
  return text.empty() || std::all_of(
    text.begin(), text.end(), [](unsigned char character) {
      return std::isspace(character) != 0;
    });
}

std::string TranslatorNode::error_code(ClientError error)
{
  switch (error) {
    case ClientError::kUnavailable:
      return "NOT_READY";
    case ClientError::kTimeout:
      return "TIMEOUT";
    case ClientError::kInvalidResponse:
      return "INVALID_RESPONSE";
    case ClientError::kNone:
    case ClientError::kUpstream:
    default:
      return "UPSTREAM_ERROR";
  }
}

}  // namespace mtran
