#include "utils.h"

#include <fstream>
#include <functional>
#include <iostream>

#include "triton/backend/backend_common.h"

namespace triton {
namespace backend {
namespace punc {
std::vector<std::vector<std::string>>
split_to_mini_sentence(const std::vector<std::string> &words,
                       size_t word_limit) {
  assert(word_limit > 1);

  std::vector<std::vector<std::string>> sentences;

  size_t length = words.size();
  if (length <= word_limit) {
    sentences.push_back(words);
    return sentences;
  }

  size_t sentence_count = length / word_limit;

  for (size_t i = 0; i < sentence_count; ++i) {
    std::vector<std::string> mini_sentence(
        words.begin() + i * word_limit, words.begin() + (i + 1) * word_limit);
    sentences.push_back(mini_sentence);
  }

  // 处理剩余部分
  if (length % word_limit > 0) {
    std::vector<std::string> last_sentence(
        words.begin() + sentence_count * word_limit, words.end());
    sentences.push_back(last_sentence);
  }

  return sentences;
}

std::vector<std::vector<int32_t>>
split_to_mini_sentence(const std::vector<int32_t> &ids, size_t id_limit) {
  std::vector<std::vector<int32_t>> mini_sentences;
  size_t i = 0;
  while (i < ids.size()) {
    size_t end = std::min(i + id_limit, ids.size());
    mini_sentences.emplace_back(ids.begin() + i, ids.begin() + end);
    i = end;
  }
  return mini_sentences;
}

std::vector<std::string> load_punc_token_list(const std::string &path) {
  std::ifstream fin(path);
  if (!fin.is_open()) {
    throw std::runtime_error("Failed to open file: " + path);
  }
  json j;
  fin >> j;
  return j.get<std::vector<std::string>>();
}

// 获取当前 UTF-8 字符的字节数
static inline size_t utf8_char_len(unsigned char c) {
  if (c < 0x80)
    return 1;
  else if ((c >> 5) == 0x6)
    return 2;
  else if ((c >> 4) == 0xE)
    return 3;
  else if ((c >> 3) == 0x1E)
    return 4;
  return 1;
}

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

  size_t i = 0;
  const size_t n = text.size();

  while (i < n) {
    unsigned char c = text[i];

    if (c == ' ') {
      if (!current_word.empty()) {
        words.push_back(current_word);
        current_word.clear();
      }
      i++;
      continue;
    }

    size_t char_len = utf8_char_len(c);
    std::string ch = text.substr(i, char_len);

    if (char_len == 1) {
      current_word += ch;
    } else {
      if (!current_word.empty()) {
        words.push_back(current_word);
        current_word.clear();
      }
      words.push_back(ch);
    }

    i += char_len;
  }

  if (!current_word.empty()) {
    words.push_back(current_word);
  }

  return words;
}

YAMLDict read_yaml(const std::string &yaml_path) {
  if (!std::filesystem::exists(yaml_path)) {
    throw std::runtime_error("YAML file does not exist: " + yaml_path);
  }

  try {
    YAML::Node data = YAML::LoadFile(yaml_path);
    return data;
  } catch (const std::exception &e) {
    throw std::runtime_error("Failed to read yaml: " + std::string(e.what()));
  }
}

std::string removeLastUtf8Char(const std::string& str) {
    if (str.empty()) {
      return str;
    }

    size_t i = str.size() - 1;
    while (i > 0 && (static_cast<unsigned char>(str[i]) & 0xC0) == 0x80) {
        --i;
    }
    return str.substr(0, i);
}

TRITONSERVER_Error *LoadModel(TRITONBACKEND_ModelInstance *instance,
                              const char *model_name, int64_t model_version) {
  TRITONBACKEND_Model *model;
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceModel(instance, &model));

  TRITONSERVER_Server *backend_server;
  TRITONSERVER_Error *err = TRITONBACKEND_ModelServer(model, &backend_server);
  if (err != nullptr) {
    std::cout << "加载模型之前有问题";
    TRITONSERVER_ErrorDelete(err);
    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL,
                                 "Cannot get server handle");
  }

  err = TRITONSERVER_ServerLoadModel(backend_server, model_name);
  if (err != nullptr) {
    std::cout << "加载模型有问题";
    std::cerr << "Failed to load model: " << model_name << std::endl;
    TRITONSERVER_ErrorDelete(err);
    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL,
                                 "Failed to load model");
  }

  std::cout << "Loaded model: " << model_name << std::endl;

  return nullptr;
}
} // namespace punc
} // namespace backend
} // namespace triton