#include "json_ffi_engine.h"

#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>

#include <filesystem>
#include <fstream>

#include "../serve/model.h"
#include "../support/json_parser.h"
#include "../support/module_vtable.h"
#include "../support/result.h"

namespace mlc {
namespace llm {
namespace json_ffi {

using namespace tvm::runtime;
using tvm::ffi::Object;
using tvm::ffi::Shape;

JSONFFIEngine::JSONFFIEngine() { engine_ = serve::ThreadedEngine::Create(); }

bool JSONFFIEngine::ChatCompletion(std::string request_json_str, std::string request_id) {
  bool success = this->AddRequest(request_json_str, request_id);
  if (!success) {
    this->StreamBackError(request_id);
  }
  return success;
}

void JSONFFIEngine::StreamBackError(std::string request_id) {
  ChatCompletionMessage delta;
  delta.content = this->err_;
  delta.role = "assistant";

  ChatCompletionStreamResponseChoice choice;
  choice.finish_reason = FinishReason::error;
  choice.index = 0;
  choice.delta = delta;

  ChatCompletionStreamResponse response;
  response.id = request_id;
  response.choices = std::vector<ChatCompletionStreamResponseChoice>{choice};
  response.model = "json_ffi";  // TODO: Return model name from engine (or from args)
  response.system_fingerprint = "";

  tvm::ffi::json::Array response_arr;
  response_arr.push_back(response.AsJSON());

  // now stream back the final usage block, which is required.
  // NOTE: always stream back final usage block as it is an
  // invariant of the system
  response.choices.clear();
  tvm::ffi::json::Object dummy_usage;
  dummy_usage.Set("prompt_tokens", static_cast<int64_t>(0));
  dummy_usage.Set("completion_tokens", static_cast<int64_t>(0));
  dummy_usage.Set("total_tokens", static_cast<int64_t>(0));
  response.usage = tvm::ffi::json::Value(dummy_usage);
  response_arr.push_back(response.AsJSON());

  std::string stream_back_json = tvm::ffi::json::Stringify(response_arr);
  this->request_stream_callback_(stream_back_json);
}

bool JSONFFIEngine::AddRequest(std::string request_json_str, std::string request_id) {
  Result<ChatCompletionRequest> request_res = ChatCompletionRequest::FromJSON(request_json_str);
  if (request_res.IsErr()) {
    err_ = request_res.UnwrapErr();
    return false;
  }
  ChatCompletionRequest request = request_res.Unwrap();
  Array<Data> inputs;
  Array<String> stop_strs;
  bool is_special_request =
      (request.debug_config.has_value() &&
       request.debug_config.value().special_request != SpecialRequestKind::kNone);
  // special request does not have to go through prompt construction
  if (!is_special_request) {
    // get prompt: note, assistant was appended in the end.
    Result<std::vector<Data>> inputs_obj =
        CreatePrompt(this->conv_template_, request, this->model_config_, this->device_);
    if (inputs_obj.IsErr()) {
      err_ = inputs_obj.UnwrapErr();
      return false;
    }
    inputs = inputs_obj.Unwrap();

    stop_strs.reserve(this->conv_template_.stop_str.size());
    for (const std::string& stop_str : this->conv_template_.stop_str) {
      stop_strs.push_back(stop_str);
    }
    if (request.stop.has_value()) {
      stop_strs.reserve(stop_strs.size() + request.stop.value().size());
      for (const std::string& stop_str : request.stop.value()) {
        stop_strs.push_back(stop_str);
      }
    }
  }
  // create a generation config from request
  const auto& default_gen_cfg = default_generation_config_;
  auto gen_cfg = tvm::ffi::make_object<GenerationConfigNode>();
  gen_cfg->n = request.n;
  gen_cfg->temperature = request.temperature.value_or(default_gen_cfg->temperature);
  gen_cfg->top_p = request.top_p.value_or(default_gen_cfg->top_p);
  gen_cfg->frequency_penalty =
      request.frequency_penalty.value_or(default_gen_cfg->frequency_penalty);
  gen_cfg->presence_penalty = request.presence_penalty.value_or(default_gen_cfg->presence_penalty);
  gen_cfg->logprobs = request.logprobs;
  gen_cfg->top_logprobs = request.top_logprobs;
  gen_cfg->logit_bias = request.logit_bias.value_or(default_gen_cfg->logit_bias);
  gen_cfg->seed = request.seed.value_or(std::random_device{}());
  gen_cfg->max_tokens = request.max_tokens.value_or(default_gen_cfg->max_tokens);
  gen_cfg->stop_strs = std::move(stop_strs);
  gen_cfg->stop_token_ids = conv_template_.stop_token_ids;
  gen_cfg->response_format = request.response_format.value_or(ResponseFormat());
  gen_cfg->debug_config = request.debug_config.value_or(DebugConfig());

  Result<GenerationConfig> res_gen_config = GenerationConfig::Validate(GenerationConfig(gen_cfg));
  if (res_gen_config.IsErr()) {
    err_ = res_gen_config.UnwrapErr();
    return false;
  }

  Request engine_request(request_id, inputs, res_gen_config.Unwrap());

  // setup request state
  RequestState rstate;
  rstate.model = request.model.value_or("");
  rstate.streamer.reserve(gen_cfg->n);
  for (int i = 0; i < gen_cfg->n; ++i) {
    rstate.streamer.push_back(TextStreamer(tokenizer_));
  }
  request_map_[request_id] = std::move(rstate);

  this->engine_->AddRequest(engine_request);
  return true;
}

bool JSONFFIEngine::Abort(std::string request_id) {
  this->engine_->AbortRequest(request_id);
  auto it = request_map_.find(request_id);
  if (it != request_map_.end()) {
    request_map_.erase(it);
  }
  return true;
}

std::string JSONFFIEngine::GetLastError() { return err_; }

void JSONFFIEngine::ExitBackgroundLoop() { this->engine_->ExitBackgroundLoop(); }

JSONFFIEngine::~JSONFFIEngine() { this->ExitBackgroundLoop(); }

class JSONFFIEngineImpl : public JSONFFIEngine, public ffi::ModuleObj {
 public:
  TVM_MODULE_VTABLE_BEGIN("mlc.json_ffi");
  TVM_MODULE_VTABLE_ENTRY("init_background_engine", &JSONFFIEngineImpl::InitBackgroundEngine);
  TVM_MODULE_VTABLE_ENTRY("reload", &JSONFFIEngineImpl::Reload);
  TVM_MODULE_VTABLE_ENTRY("unload", &JSONFFIEngineImpl::Unload);
  TVM_MODULE_VTABLE_ENTRY("reset", &JSONFFIEngineImpl::Reset);
  TVM_MODULE_VTABLE_ENTRY("chat_completion", &JSONFFIEngineImpl::ChatCompletion);
  TVM_MODULE_VTABLE_ENTRY("abort", &JSONFFIEngineImpl::Abort);
  TVM_MODULE_VTABLE_ENTRY("get_last_error", &JSONFFIEngineImpl::GetLastError);
  TVM_MODULE_VTABLE_ENTRY("run_background_loop", &JSONFFIEngineImpl::RunBackgroundLoop);
  TVM_MODULE_VTABLE_ENTRY("run_background_stream_back_loop",
                          &JSONFFIEngineImpl::RunBackgroundStreamBackLoop);
  TVM_MODULE_VTABLE_ENTRY("exit_background_loop", &JSONFFIEngineImpl::ExitBackgroundLoop);
  TVM_MODULE_VTABLE_END();

  void InitBackgroundEngine(int device_type, int device_id,
                            Optional<Function> request_stream_callback) {
    DLDevice device{static_cast<DLDeviceType>(device_type), device_id};
    this->device_ = device;
    TVM_FFI_ICHECK(request_stream_callback.has_value())
        << "JSONFFIEngine requires request stream callback function, but it is not given.";
    this->request_stream_callback_ = request_stream_callback.value();

    auto frequest_stream_callback_wrapper = [this](ffi::PackedArgs args, ffi::Any* ret) {
      TVM_FFI_ICHECK_EQ(args.size(), 1);
      Array<RequestStreamOutput> delta_outputs = args[0].cast<Array<RequestStreamOutput>>();
      std::string responses = this->GetResponseFromStreamOutput(delta_outputs);
      this->request_stream_callback_(responses);
    };

    request_stream_callback = Function(frequest_stream_callback_wrapper);
    this->engine_->InitThreadedEngine(device, std::move(request_stream_callback), std::nullopt);
  }

  void Reload(String engine_config_json_str) {
    this->engine_->Reload(engine_config_json_str);
    this->default_generation_config_ = this->engine_->GetDefaultGenerationConfig();
    auto engine_config = this->engine_->GetCompleteEngineConfig();

    // Load conversation template.
    Result<tvm::ffi::json::Object> model_config_json =
        serve::Model::LoadModelConfig(engine_config->model);
    TVM_FFI_ICHECK(model_config_json.IsOk()) << model_config_json.UnwrapErr();
    const tvm::ffi::json::Object& model_config_json_unwrapped = model_config_json.Unwrap();
    Result<Conversation> conv_template = Conversation::FromJSON(
        json::Lookup<tvm::ffi::json::Object>(model_config_json_unwrapped, "conv_template"));
    TVM_FFI_ICHECK(!conv_template.IsErr())
        << "Invalid conversation template JSON: " << conv_template.UnwrapErr();
    this->conv_template_ = conv_template.Unwrap();
    this->model_config_ = ModelConfig::FromJSON(
        json::Lookup<tvm::ffi::json::Object>(model_config_json_unwrapped, "model_config"));
    this->tokenizer_ = Tokenizer::FromPath(engine_config->model);
  }

  void Unload() { this->engine_->Unload(); }

  void Reset() { this->engine_->Reset(); }

  void RunBackgroundLoop() { this->engine_->RunBackgroundLoop(); }

  void RunBackgroundStreamBackLoop() { this->engine_->RunBackgroundStreamBackLoop(); }

  String GetResponseFromStreamOutput(Array<RequestStreamOutput> delta_outputs) {
    tvm::ffi::json::Array json_response_arr;
    for (const auto& delta_output : delta_outputs) {
      std::string request_id = delta_output->request_id;
      auto request_state_it = request_map_.find(request_id);
      if (request_state_it == request_map_.end()) continue;
      RequestState& rstate = request_state_it->second;

      // build the final usage messages
      // invariant, we can always let other messages to come first
      // then the final usage messages, as final usage is always last
      if (delta_output->request_final_usage_json_str.has_value()) {
        ChatCompletionStreamResponse response;
        response.id = request_id;
        response.model = rstate.model;
        response.system_fingerprint = "";
        std::string usage_json_str = delta_output->request_final_usage_json_str.value();
        tvm::ffi::String parse_err;
        auto usage_json = tvm::ffi::json::Parse(usage_json_str, &parse_err);
        if (!parse_err.empty()) {
          err_ = parse_err;
        } else {
          response.usage = usage_json;
        }
        json_response_arr.push_back(response.AsJSON());
        request_map_.erase(request_state_it);
        continue;
      }
      TVM_FFI_ICHECK_NE(delta_output->group_finish_reason.size(), 0);
      TVM_FFI_ICHECK_EQ(delta_output->group_delta_token_ids.size(),
                        delta_output->group_finish_reason.size());
      TVM_FFI_ICHECK_EQ(delta_output->group_delta_token_ids.size(), rstate.streamer.size());

      ChatCompletionStreamResponse response;
      response.id = request_id;
      response.model = rstate.model;
      response.system_fingerprint = "";

      for (size_t i = 0; i < delta_output->group_finish_reason.size(); ++i) {
        // choice
        ChatCompletionStreamResponseChoice choice;
        Optional<String> finish_reason = delta_output->group_finish_reason[i];
        if (finish_reason.has_value()) {
          if (finish_reason.value() == "stop") {
            choice.finish_reason = FinishReason::stop;
          } else if (finish_reason.value() == "length") {
            choice.finish_reason = FinishReason::length;
          } else if (finish_reason.value() == "tool_calls") {
            choice.finish_reason = FinishReason::tool_calls;
          } else if (finish_reason.value() == "error") {
            choice.finish_reason = FinishReason::error;
          }
        } else {
          choice.finish_reason = std::nullopt;
        }
        choice.index = static_cast<int>(i);
        ChatCompletionMessage delta;
        // Size of delta_output->group_delta_token_ids Array should be 1
        const Shape& delta_token_ids = delta_output->group_delta_token_ids[i];
        std::vector<int32_t> delta_token_ids_vec(delta_token_ids.begin(), delta_token_ids.end());
        std::string content = rstate.streamer[i]->Put(delta_token_ids_vec);
        if (finish_reason.has_value()) {
          content += rstate.streamer[i]->Finish();
        }
        if (!content.empty()) {
          delta.content = content;
        }
        delta.role = "assistant";
        choice.delta = delta;
        if (!choice.delta.content.IsNull() || choice.finish_reason.has_value()) {
          response.choices.push_back(choice);
        }
      }
      // if it is not the usage block, choices cannot be empty
      if (!response.choices.empty()) {
        json_response_arr.push_back(response.AsJSON());
      }
    }
    return tvm::ffi::json::Stringify(json_response_arr);
  }
};

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("mlc.json_ffi.CreateJSONFFIEngine",
                        []() { return ffi::Module(tvm::ffi::make_object<JSONFFIEngineImpl>()); });
}

}  // namespace json_ffi
}  // namespace llm
}  // namespace mlc