#include <yaml-cpp/yaml.h>

#include <nlohmann/json.hpp>
#include <string>
#include <vector>

#include "triton/core/tritonbackend.h"
#include "triton/core/tritonserver.h"

namespace triton {
namespace backend {
namespace punc {
using json = nlohmann::json;
using YAMLDict = YAML::Node;

std::vector<std::vector<std::string>>
split_to_mini_sentence(const std::vector<std::string> &words,
                       size_t word_limit = 20);
std::vector<std::vector<int32_t>>
split_to_mini_sentence(const std::vector<int32_t> &ids, size_t id_limit = 20);

std::vector<std::string> load_punc_token_list(const std::string &path);

std::vector<std::string> code_mix_split_words(const std::string &text);

YAMLDict read_yaml(const std::string &yaml_path);

std::string removeLastUtf8Char(const std::string& str);

TRITONSERVER_Error *
LoadModel(TRITONBACKEND_ModelInstance *instance,
          const char *model_name = "punc_ct-transformer_nonstreaming",
          int64_t model_version = -1);

class TokenIDConverter {
public:
  TokenIDConverter(const std::vector<std::string> &token_list)
      : token_list_(token_list) {
    if (token_list_.empty()) {
      throw std::runtime_error("token_list cannot be empty");
    }

    unk_symbol_ = token_list_.back();
    // 构建 token -> id 映射
    for (size_t i = 0; i < token_list_.size(); ++i) {
      token2id_[token_list_[i]] = static_cast<int>(i);
    }
    unk_id_ = token2id_[unk_symbol_];
  }

  int get_num_vocabulary_size() const {
    return static_cast<int>(token_list_.size());
  }

  std::vector<std::string> ids2tokens(const std::vector<int> &integers) const {
    std::vector<std::string> result;
    result.reserve(integers.size());
    for (int i : integers) {
      if (i >= 0 && i < static_cast<int>(token_list_.size()))
        result.push_back(token_list_[i]);
      else
        result.push_back(unk_symbol_);
    }
    return result;
  }

  std::vector<int> tokens2ids(const std::vector<std::string> &tokens) const {
    std::vector<int> result;
    result.reserve(tokens.size());
    for (const auto &t : tokens) {
      auto it = token2id_.find(t);
      if (it != token2id_.end())
        result.push_back(it->second);
      else
        result.push_back(unk_id_);
    }
    return result;
  }

private:
  std::vector<std::string> token_list_;
  std::string unk_symbol_;
  int unk_id_;
  std::unordered_map<std::string, int> token2id_;
};

} // namespace punc
} // namespace backend
} // namespace triton