#include "triton/backend/backend_common.h"
#include "triton/backend/backend_input_collector.h"
#include "triton/backend/backend_model.h"
#include "triton/backend/backend_model_instance.h"
#include "triton/backend/backend_output_responder.h"
#include "triton/core/tritonbackend.h"
#include "triton/core/tritonserver.h"
#include "utils.h"

#include <unistd.h>

#include <atomic>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstring>
#include <fstream>
#include <iostream>
#include <memory>
#include <string>

namespace triton {
namespace backend {
namespace punc {
struct PuncCallbackData {
  std::function<void(std::vector<int>, TRITONSERVER_Error *)> user_callback;
  TRITONSERVER_ResponseAllocator *allocator;
  TRITONSERVER_InferenceRequest *request;
  bool called = false;
  // 存储 text_lengths 数据，确保其生命周期与请求一致
  std::shared_ptr<int32_t> text_lengths_data;
};

TRITONSERVER_Error *ResponseAlloc(TRITONSERVER_ResponseAllocator *allocator,
                                  const char *tensor_name, size_t byte_size,
                                  TRITONSERVER_MemoryType preferred_memory_type,
                                  int64_t preferred_memory_type_id, void *userp,
                                  void **buffer, void **buffer_userp,
                                  TRITONSERVER_MemoryType *actual_memory_type,
                                  int64_t *actual_memory_type_id) {
  *buffer = malloc(byte_size);
  *buffer_userp = nullptr;
  *actual_memory_type = TRITONSERVER_MEMORY_CPU;
  *actual_memory_type_id = 0;
  return nullptr; // success
}

TRITONSERVER_Error *ResponseRelease(TRITONSERVER_ResponseAllocator *allocator,
                                    void *buffer, void *buffer_userp,
                                    size_t byte_size,
                                    TRITONSERVER_MemoryType memory_type,
                                    int64_t memory_type_id) {
  free(buffer);
  return nullptr;
}

void PuncResponseCompleteCallback(TRITONSERVER_InferenceResponse *response,
                                  const uint32_t flags, void *userp) {
  auto *data = static_cast<PuncCallbackData *>(userp);
  if (data->called)
    return;
  data->called = true;

  std::vector<int> punct_indices;
  TRITONSERVER_Error *err = nullptr;

  if (response == nullptr) {
    err = TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, "Null response");
  } else {
    const char *output_name;
    TRITONSERVER_DataType dtype;
    const int64_t *shape;
    uint64_t dim_count;
    const void *base;
    size_t byte_size;
    TRITONSERVER_MemoryType memory_type;
    int64_t memory_type_id;
    void *userdata;

    err = TRITONSERVER_InferenceResponseOutput(
        response, 0, &output_name, &dtype, &shape, &dim_count, &base,
        &byte_size, &memory_type, &memory_type_id, &userdata);

    if (!err) {
      if (dtype != TRITONSERVER_TYPE_FP32) {
        err = TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INVALID_ARG,
                                    "Expected FP32 logits");
      } else if (dim_count != 3 || shape[0] != 1 || shape[2] != 6) {
        err = TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INVALID_ARG,
                                    "Logits shape must be [1, seq_len, 6]");
      }
    }

    if (!err) {
      size_t seq_len = shape[1];
      const float *logits = static_cast<const float *>(base);
      punct_indices.reserve(seq_len);
      for (size_t i = 0; i < seq_len; ++i) {
        const float *pos = logits + i * 6;
        int best = 0;
        for (int k = 1; k < 6; ++k) {
          if (pos[k] > pos[best])
            best = k;
        }
        if (best < 0 || best >= 6)
          best = 0;
        punct_indices.push_back(best);
      }
    }
  }

  if (punct_indices.empty() == true) {
    std::cout << "回调：No punctuation indices found." << std::endl;
  }

  data->user_callback(std::move(punct_indices), err);

  if (response != nullptr) {
    TRITONSERVER_InferenceResponseDelete(response);
  }
}

void InferRequestCompleteCallback(TRITONSERVER_InferenceRequest *request,
                                  uint32_t flags, void *userp) {
  auto *data = reinterpret_cast<PuncCallbackData *>(userp);

  if (data == nullptr) {
    return;
  }

  if (data->allocator != nullptr) {
    TRITONSERVER_ResponseAllocatorDelete(data->allocator);
    data->allocator = nullptr;
  }

  if (request != nullptr) {
    TRITONSERVER_InferenceRequestDelete(request);
    request = nullptr;
  }

  delete data;
  userp = nullptr;
}

TRITONSERVER_Error *RunPuncModelAsync(
    TRITONBACKEND_ModelInstance *instance,
    const std::vector<int32_t> &input_ids,
    std::function<void(std::vector<int>, TRITONSERVER_Error *)> callback) {
  TRITONBACKEND_Model *model;
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceModel(instance, &model));

  TRITONSERVER_Server *server;
  RETURN_IF_ERROR(TRITONBACKEND_ModelServer(model, &server));

  TRITONSERVER_InferenceRequest *infer_request = nullptr;
  const char *target_model = "punc_ct-transformer_nonstreaming";
  bool is_ready = false;
  RETURN_IF_ERROR(
      TRITONSERVER_ServerModelIsReady(server, target_model, -1, &is_ready));

  if (!is_ready) {
    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_UNAVAILABLE,
                                 "Model is not ready");
  }
  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestNew(&infer_request, server,
                                                   target_model, -1));

  TRITONSERVER_ResponseAllocator *allocator = nullptr;
  RETURN_IF_ERROR(TRITONSERVER_ResponseAllocatorNew(&allocator, ResponseAlloc,
                                                    ResponseRelease, nullptr));

  auto *cb_data =
      new PuncCallbackData{std::move(callback), allocator, infer_request, false, nullptr};
  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestSetResponseCallback(
      infer_request, allocator, nullptr, PuncResponseCompleteCallback,
      cb_data));
  TRITONSERVER_InferenceRequestSetReleaseCallback(
      infer_request, InferRequestCompleteCallback, cb_data);

  int64_t shape_inputs[2] = {1, static_cast<int64_t>(input_ids.size())};
  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestAddInput(
      infer_request, "inputs", TRITONSERVER_TYPE_INT32, shape_inputs, 2));
  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestAppendInputData(
      infer_request, "inputs", input_ids.data(),
      input_ids.size() * sizeof(int32_t), TRITONSERVER_MEMORY_CPU, 0));

  // 使用 shared_ptr 保持 text_lengths 数据的生命周期
  auto len_val_ptr =
      std::make_shared<int32_t>(static_cast<int32_t>(input_ids.size()));
  int64_t shape_len[2] = {1, 1};
  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestAddInput(
      infer_request, "text_lengths", TRITONSERVER_TYPE_INT32, shape_len, 2));

  // 将 shared_ptr 存储到 callback data 中，确保在请求完成前不被释放
  cb_data->text_lengths_data = len_val_ptr;

  RETURN_IF_ERROR(TRITONSERVER_InferenceRequestAppendInputData(
      infer_request, "text_lengths", len_val_ptr.get(), sizeof(int32_t),
      TRITONSERVER_MEMORY_CPU, 0));

  RETURN_IF_ERROR(
      TRITONSERVER_InferenceRequestAddRequestedOutput(infer_request, "logits"));

  TRITONSERVER_Error *err =
      TRITONSERVER_ServerInferAsync(server, infer_request, nullptr);

  if (err != nullptr) {
    std::cout << "异步推理失败，错误码：" << TRITONSERVER_ErrorCode(err)
              << "错误信息：" << TRITONSERVER_ErrorMessage(err) << std::endl;
    delete cb_data;
    TRITONSERVER_InferenceRequestDelete(infer_request);
    TRITONSERVER_ResponseAllocatorDelete(allocator);
    return err;
  }

  return nullptr;
}

TRITONSERVER_Error *RunPuncModelSync(TRITONBACKEND_ModelInstance *instance,
                                     const std::vector<int32_t> &input_ids,
                                     std::vector<int> &punct_indices) {
  // 使用 shared_ptr 管理回调数据，确保生命周期安全
  struct SyncCallbackData
  {
    std::mutex mtx;
    std::condition_variable cv;
    bool finished = false;
    TRITONSERVER_Error *async_err = nullptr;
    std::vector<int> result;
  };

  auto sync_data = std::make_shared<SyncCallbackData>();

  auto callback = [sync_data](std::vector<int> indices, TRITONSERVER_Error *err)
  {
    {
      std::unique_lock<std::mutex> lock(sync_data->mtx);
      sync_data->result = std::move(indices);
      sync_data->async_err = err;
      sync_data->finished = true;
    }
    sync_data->cv.notify_one();
  };

  TRITONSERVER_Error *submit_err =
      RunPuncModelAsync(instance, input_ids, callback);
  if (submit_err != nullptr) {
    return submit_err;
  }

  std::unique_lock<std::mutex> lock(sync_data->mtx);

  while (!sync_data->finished)
  {
    sync_data->cv.wait_for(lock, std::chrono::milliseconds(100));
  }

  punct_indices = std::move(sync_data->result);

  if (sync_data->async_err != nullptr)
  {
    return sync_data->async_err;
  }

  if (punct_indices.empty())
  {
    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL,
                                 "Empty result from punctuation model");
  }

  return nullptr;
}

class ModelState : public BackendModel {
public:
  static TRITONSERVER_Error *Create(TRITONBACKEND_Model *triton_model,
                                    ModelState **state);
  virtual ~ModelState() = default;

private:
  ModelState(TRITONBACKEND_Model *triton_model) : BackendModel(triton_model) {}
};

TRITONSERVER_Error *ModelState::Create(TRITONBACKEND_Model *triton_model,
                                       ModelState **state) {
  try {
    *state = new ModelState(triton_model);
  } catch (const BackendModelException &ex) {
    RETURN_ERROR_IF_TRUE(
        ex.err_ == nullptr, TRITONSERVER_ERROR_INTERNAL,
        std::string("unexpected nullptr in BackendModelException"));
    RETURN_IF_ERROR(ex.err_);
  }

  return nullptr;
}

extern "C" {
TRITONSERVER_Error *TRITONBACKEND_ModelInitialize(TRITONBACKEND_Model *model) {
  ModelState *model_state;
  RETURN_IF_ERROR(ModelState::Create(model, &model_state));
  RETURN_IF_ERROR(TRITONBACKEND_ModelSetState(
      model, reinterpret_cast<void *>(model_state)));

  return nullptr;
}

TRITONSERVER_Error *TRITONBACKEND_ModelFinalize(TRITONBACKEND_Model *model) {
  void *vstate;
  RETURN_IF_ERROR(TRITONBACKEND_ModelState(model, &vstate));
  ModelState *model_state = reinterpret_cast<ModelState *>(vstate);
  delete model_state;

  return nullptr;
}
}

class ModelInstanceState : public BackendModelInstance {
public:
  static TRITONSERVER_Error *
  Create(ModelState *model_state,
         TRITONBACKEND_ModelInstance *triton_model_instance,
         ModelInstanceState **state);
  virtual ~ModelInstanceState() = default;

  // Get the state of the model that corresponds to this instance.
  ModelState *StateForModel() const { return model_state_; }

  std::vector<std::string> punc_list_;
  int period_index_ = 0;
  std::shared_ptr<TokenIDConverter> converter_;

private:
  ModelInstanceState(ModelState *model_state,
                     TRITONBACKEND_ModelInstance *triton_model_instance)
      : BackendModelInstance(model_state, triton_model_instance),
        model_state_(model_state) {}

  ModelState *model_state_;
};

TRITONSERVER_Error *
ModelInstanceState::Create(ModelState *model_state,
                           TRITONBACKEND_ModelInstance *triton_model_instance,
                           ModelInstanceState **state) {
  try {
    *state = new ModelInstanceState(model_state, triton_model_instance);
  } catch (const BackendModelInstanceException &ex) {
    RETURN_ERROR_IF_TRUE(
        ex.err_ == nullptr, TRITONSERVER_ERROR_INTERNAL,
        std::string("unexpected nullptr in BackendModelInstanceException"));
    RETURN_IF_ERROR(ex.err_);
  }

  return nullptr;
}

extern "C" {
TRITONSERVER_Error *
TRITONBACKEND_ModelInstanceInitialize(TRITONBACKEND_ModelInstance *instance) {
  TRITONBACKEND_Model *model;
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceModel(instance, &model));

  void *vmodelstate;
  RETURN_IF_ERROR(TRITONBACKEND_ModelState(model, &vmodelstate));
  ModelState *model_state = reinterpret_cast<ModelState *>(vmodelstate);

  ModelInstanceState *instance_state;
  RETURN_IF_ERROR(
      ModelInstanceState::Create(model_state, instance, &instance_state));
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceSetState(
      instance, reinterpret_cast<void *>(instance_state)));

  TRITONBACKEND_ArtifactType artifact_type;
  const char *repository_path = nullptr;
  auto err =
      TRITONBACKEND_ModelRepository(model, &artifact_type, &repository_path);
  if (err != nullptr) {
    std::cout << "获取模型路径失败，错误码：" << TRITONSERVER_ErrorCode(err)
              << "，错误信息：" << TRITONSERVER_ErrorMessage(err) << std::endl;
    return err;
  }
  std::string punc_config_path =
      std::string(repository_path) + "/1/config.yaml";
  std::string punc_token_list_path =
      std::string(repository_path) + "/1/tokens_punc.json";

  std::vector<std::string> punc_token_list =
      load_punc_token_list(punc_token_list_path);

  instance_state->converter_ =
      std::make_shared<TokenIDConverter>(punc_token_list);

  YAMLDict punc_config = read_yaml(punc_config_path);

  auto punc_list_node = punc_config["model_conf"]["punc_list"];

  instance_state->punc_list_.clear();
  for (size_t i = 0; i < punc_list_node.size(); ++i) {
    std::string punc = punc_list_node[i].as<std::string>();
    if (punc == ",")
      punc = "，";
    else if (punc == "?")
      punc = "？";
    else if (punc == "。")
      instance_state->period_index_ = i;
    instance_state->punc_list_.push_back(punc);
  }

  auto loadModelResult = LoadModel(instance);
  if (loadModelResult != nullptr) {
    std::cout << "加载模型失败" << std::endl;
    return loadModelResult;
  }

  return nullptr; // success
}

TRITONSERVER_Error *
TRITONBACKEND_ModelInstanceFinalize(TRITONBACKEND_ModelInstance *instance)
{
  void *vstate;
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceState(instance, &vstate));
  ModelInstanceState *instance_state =
      reinterpret_cast<ModelInstanceState *>(vstate);
  delete instance_state;

  return nullptr;
}

} // extern "C"

extern "C" {

TRITONSERVER_Error *
TRITONBACKEND_ModelInstanceExecute(TRITONBACKEND_ModelInstance *instance,
                                   TRITONBACKEND_Request **requests,
                                   const uint32_t request_count) {
  ModelInstanceState *instance_state;
  RETURN_IF_ERROR(TRITONBACKEND_ModelInstanceState(
      instance, reinterpret_cast<void **>(&instance_state)));

  int split_size = 20;
  for (uint32_t r = 0; r < request_count; ++r) {
    auto &request = requests[r];
    // 获取输入
    TRITONBACKEND_Input *input = nullptr;
    TRITONSERVER_Error *err =
        TRITONBACKEND_RequestInput(request, "text_no_punc", &input);
    RETURN_IF_ERROR(err);
    const int64_t *shape;
    uint32_t dims_count;
    RETURN_IF_ERROR(TRITONBACKEND_InputProperties(
        input, nullptr, nullptr, &shape, &dims_count, nullptr, nullptr));
    const void *buffer;
    size_t byte_size;
    TRITONSERVER_MemoryType mem_type;
    int64_t mem_id;

    RETURN_IF_ERROR(TRITONBACKEND_InputBuffer(input, 0, &buffer, &byte_size,
                                              &mem_type, &mem_id));

    std::vector<std::string> text_list;

    size_t offset = 0;
    while (offset < byte_size) {
      uint32_t len;
      memcpy(&len, (char *)buffer + offset, sizeof(uint32_t));
      offset += sizeof(uint32_t);

      std::string s((char *)buffer + offset, len);
      offset += len;

      text_list.push_back(s);
    }

    std::string all_text_inputs;
    for (const auto &s : text_list) {
      all_text_inputs += s;
    }

    std::vector<std::string> split_text = code_mix_split_words(all_text_inputs);
    std::vector<int32_t> split_text_tokens =
        instance_state->converter_->tokens2ids(split_text);
    auto mini_sentences = split_to_mini_sentence(split_text, split_size);
    auto mini_sentences_id =
        split_to_mini_sentence(split_text_tokens, split_size);

    assert(mini_sentences.size() == mini_sentences_id.size());

    std::string new_mini_sentence;
    std::vector<int> new_mini_sentence_punc;
    std::vector<int> punct_indices;
    std::vector<std::string> mini_sentence;
    std::vector<int32_t> mini_sentence_id;
    size_t cache_pop_trigger_limit = 200;
    std::vector<int32_t> cache_sent_id;
    std::vector<std::string> cache_sent;
    std::string sentence_out;
    for (size_t mini_sentence_i = 0; mini_sentence_i < mini_sentences.size();
         ++mini_sentence_i) {
      mini_sentence = mini_sentences[mini_sentence_i];
      mini_sentence.insert(mini_sentence.begin(), cache_sent.begin(),
                           cache_sent.end());
      mini_sentence_id = mini_sentences_id[mini_sentence_i];
      mini_sentence_id.insert(mini_sentence_id.begin(), cache_sent_id.begin(),
                              cache_sent_id.end());

      TRITONSERVER_Error *err =
          RunPuncModelSync(instance, mini_sentence_id, punct_indices);
      if (err) {
        std::cerr << "Punc model error: "
                  << (TRITONSERVER_ErrorMessage(err)
                          ? TRITONSERVER_ErrorMessage(err)
                          : "unknown")
                  << std::endl;
        TRITONSERVER_ErrorDelete(err);
        cache_sent.clear();
        cache_sent_id.clear();
        continue;
      } else {
        if (punct_indices.empty() == true) {
          std::cout << "No punctuation indices found." << std::endl;
        }
      }

      assert(punct_indices.size() == mini_sentence.size());
      if (mini_sentence_i < mini_sentences.size() - 1) {
        int sentenceEnd = -1;
        int last_comma_index = -1;

        for (int i = punct_indices.size() - 2; i > 1; --i) {
          if (instance_state->punc_list_[punct_indices[i]] == "。" ||
              instance_state->punc_list_[punct_indices[i]] == "？") {
            sentenceEnd = i;
            break;
          }

          if (last_comma_index < 0 &&
              instance_state->punc_list_[punct_indices[i]] == "，") {
            last_comma_index = i;
          }
        }
        if (sentenceEnd < 0 && mini_sentence.size() > cache_pop_trigger_limit &&
            last_comma_index >= 0) {
          sentenceEnd = last_comma_index;
          punct_indices[sentenceEnd] = instance_state->period_index_;
        }
        if (sentenceEnd < 0) {
          cache_sent.clear();
          cache_sent_id.clear();
        }

        cache_sent.assign(mini_sentence.begin() + sentenceEnd + 1,
                          mini_sentence.end());

        cache_sent_id.assign(mini_sentence_id.begin() + sentenceEnd + 1,
                             mini_sentence_id.end());

        mini_sentence.resize(sentenceEnd + 1);
        punct_indices.resize(sentenceEnd + 1);
      }
      new_mini_sentence_punc.insert(new_mini_sentence_punc.end(),
                                    punct_indices.begin(), punct_indices.end());
      std::vector<std::string> words_with_punc;

      for (size_t i = 0; i < mini_sentence.size(); i++) {
        if (i > 0) {
          if (mini_sentence[i].size() == 1 &&
              mini_sentence[i - 1].size() == 1) {
            mini_sentence[i] = " " + mini_sentences[0][i];
          }
        }
        words_with_punc.push_back(mini_sentence[i]);

        if (instance_state->punc_list_[punct_indices[i]] != "_") {
          words_with_punc.push_back(
              instance_state->punc_list_[punct_indices[i]]);
        }
      }
      for (const auto &w : words_with_punc) {
        new_mini_sentence += w;
      }

      new_mini_sentence_punc.insert(new_mini_sentence_punc.end(),
                                    punct_indices.begin(), punct_indices.end());

      std::string new_mini_sentence_out = new_mini_sentence;
      std::vector<int> new_mini_sentence_punc_out = new_mini_sentence_punc;

      if (mini_sentence_i == mini_sentences.size() - 1) {
        if (!new_mini_sentence.empty()) {
          size_t len = new_mini_sentence.size();
          std::string last_char = (len >= 3)
                                      ? new_mini_sentence.substr(len - 3, 3)
                                      : new_mini_sentence;

          if (last_char == "，" || last_char == "、") {
            new_mini_sentence = removeLastUtf8Char(new_mini_sentence);
            new_mini_sentence_out=new_mini_sentence+"。";
            new_mini_sentence_punc_out.back() = instance_state->period_index_;
          } else if (last_char != "。" && last_char != "？") {
            new_mini_sentence_out += "。";
            new_mini_sentence_punc_out.back() = instance_state->period_index_;
          }
        }
      }
      sentence_out = new_mini_sentence_out;
    }

    TRITONBACKEND_Response *response = nullptr;
    RETURN_IF_ERROR(TRITONBACKEND_ResponseNew(&response, request));

    TRITONBACKEND_Output *output = nullptr;
    int64_t output_shape[1] = {1};
    RETURN_IF_ERROR(TRITONBACKEND_ResponseOutput(response, &output, "OUTPUT0",
                                                 TRITONSERVER_TYPE_BYTES,
                                                 output_shape, 1));

    void *buffer_out = nullptr;
    size_t buffer_size = sizeof(uint32_t) + sentence_out.size();
    TRITONSERVER_MemoryType mem_type_out = TRITONSERVER_MEMORY_CPU;
    int64_t mem_type_id = 0;
    RETURN_IF_ERROR(TRITONBACKEND_OutputBuffer(output, &buffer_out, buffer_size,
                                               &mem_type_out, &mem_type_id));

    if (sentence_out.size() > buffer_size) {
      return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL,
                                   "Output buffer too small for sentence_out");
    }

    uint32_t str_len = static_cast<uint32_t>(sentence_out.size());
    memcpy(buffer_out, &str_len, sizeof(uint32_t));

    memcpy((char *)buffer_out + sizeof(uint32_t), sentence_out.data(), str_len);

    RETURN_IF_ERROR(TRITONBACKEND_ResponseSend(
        response, TRITONSERVER_RESPONSE_COMPLETE_FINAL, nullptr));
  }

  return nullptr; // success
}

} // extern "C"

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