/*!
 *  Copyright (c) 2023-2025 by Contributors
 * \file serve/data.h
 */
#ifndef MLC_LLM_SERVE_DATA_H_
#define MLC_LLM_SERVE_DATA_H_

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/runtime/tensor.h>

#include <atomic>
#include <optional>

#include "../tokenizers/tokenizers.h"

namespace mlc {
namespace llm {
namespace serve {

using namespace tvm::runtime;
using tvm::ffi::Object;
using tvm::ffi::ObjectPtr;
using tvm::ffi::ObjectRef;
using tvm::ffi::Optional;
using tvm::ffi::Shape;

class Model;

/****************** DataNode ******************/

/*! \brief The base class of multi-modality data (text, tokens, embedding, etc). */
class DataNode : public Object {
 public:
  /*! \brief Get the length (equivalent number of tokens) of the data. */
  virtual int GetLength() const = 0;

  /*!
   * \brief Compute the embedding of this data with regard to the input model.
   * When the input destination pointer is not nullptr, it in-place writes the
   * embedding into the input destination array at the given offset.
   * Otherwise, the embeddings will be directly returned back.
   * \param model The model to take embeddings from.
   * \param dst The destination array of the embedding lookup.
   * \param offset The token offset where the computed embeddings will be written
   * into the destination array.
   * \return The updated destination embedding array or the computed embeddings.
   * \note When `dst` is nullptr, we require `offset` to be 0.
   */
  virtual ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const = 0;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataNode>();
  }

  static constexpr const bool _type_has_method_sequal_reduce = false;
  static constexpr const bool _type_has_method_shash_reduce = false;
  static constexpr const uint32_t _type_child_slots = 3;
  TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.Data", DataNode, Object);
};

class Data : public ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Data, ObjectRef, DataNode);
};

/*! \brief Split the given data array into two arrays at the "split_pos" position. */
std::pair<Array<Data>, Array<Data>> SplitData(const Array<Data>& original_data, int total_length,
                                              int split_pos);

/****************** TextDataNode ******************/

/*! \brief The class of text data, containing a text string. */
class TextDataNode : public DataNode {
 public:
  /*! \brief The text string. */
  tvm::ffi::String text;

  int GetLength() const final;
  ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TextDataNode>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TextData", TextDataNode, DataNode);
};

class TextData : public Data {
 public:
  explicit TextData(String text);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TextData, Data, TextDataNode);
};

/****************** TokenDataNode ******************/

/*! \brief The class of token data, containing a list of token ids. */
class TokenDataNode : public DataNode {
 public:
  /*! \brief The token ids. */
  Shape token_ids;

  int GetLength() const final;
  ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TokenDataNode>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.TokenData", TokenDataNode, DataNode);
};

class TokenData : public Data {
 public:
  explicit TokenData(Shape token_ids);

  explicit TokenData(std::vector<int32_t> token_ids);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TokenData, Data, TokenDataNode);
};

/****************** ImageDataNode ******************/

/*! \brief The class of image data, containing a 3D array of pixel values. */
class ImageDataNode : public DataNode {
 public:
  /*! \brief The pixel values. */
  Tensor image;
  int embed_size;

  int GetLength() const final;
  ObjectRef GetEmbedding(Model model, ObjectRef* dst = nullptr, int offset = 0) const final;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ImageDataNode>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.serve.ImageData", ImageDataNode, DataNode);
};

class ImageData : public Data {
 public:
  explicit ImageData(Tensor image, int embed_size);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ImageData, Data, ImageDataNode);
};

/****************** SampleResult ******************/

// The pair of a token id and its probability in sampling.
using TokenProbPair = std::pair<int32_t, float>;

/*!
 * \brief The class of sampler's sampling result.
 * It's not a TVM object since it will not be used directly on Python side.
 */
struct SampleResult {
  /*! \brief The token id and probability of the sampled token. */
  TokenProbPair sampled_token_id;
  /*! \brief The token id and probability of the tokens with top probabilities. */
  std::vector<TokenProbPair> top_prob_tokens;

  /*! \brief Get the sampled token id. */
  int32_t GetTokenId() const;

  /*!
   * \brief Get the logprob JSON string of this token with regard
   * to OpenAI API at https://platform.openai.com/docs/api-reference/chat/object.
   * \param tokenizer The tokenizer for token table lookup.
   * \param logprob A boolean indicating if need to return log probability.
   * \return A JSON string that conforms to the logprob spec in OpenAI API.
   */
  std::string GetLogProbJSON(const Tokenizer& tokenizer, bool logprob) const;
};

/****************** RequestStreamOutput ******************/

/*!
 * \brief The generated delta request output that is streamed back
 * through callback stream function.
 *
 * \note: This output object corresponds to parallel generated outputs when n != 1.
 *
 * For example, if n=2, then group_delta_token_ids[0] matches to the output stream 0
 * and group_delta_token_ids[1] matches to the output stream 1
 */
class RequestStreamOutputObj : public Object {
 public:
  /*! \brief The id of the request that the function is invoked for. */
  String request_id;
  /*!
   * \brief The new generated token ids since the last callback invocation
   * for the input request.
   */
  std::vector<std::vector<int64_t>> group_delta_token_ids;
  /*! \brief The logprobs JSON strings of the new generated tokens since last invocation. */
  std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs;
  /*!
   * \brief The finish reason of the request when it is finished,
   * of None if the request has not finished yet.
   */
  std::vector<Optional<String>> group_finish_reason;
  /*!
   * \brief The usage field of the response, this is global to all streams.
   */
  Optional<String> request_final_usage_json_str;

  /*!
   * \brief The extra prefix string of all requests.
   */
  std::vector<String> group_extra_prefix_string;

  std::atomic<bool> unpacked = false;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RequestStreamOutputObj>();
  }

  static constexpr const bool _type_has_method_sequal_reduce = false;
  static constexpr const bool _type_has_method_shash_reduce = false;
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.RequestStreamOutput", RequestStreamOutputObj, Object);
};

/*!
 * \brief Managed reference to RequestStreamOutputObj.
 * \sa RequestStreamOutputObj
 */
class RequestStreamOutput : public ObjectRef {
 public:
  explicit RequestStreamOutput(
      String request_id, std::vector<std::vector<int64_t>> group_delta_token_ids,
      std::optional<std::vector<std::vector<String>>> group_delta_logprob_json_strs,
      std::vector<Optional<String>> group_finish_reason,
      std::vector<String> group_extra_prefix_string);

  static RequestStreamOutput Usage(String request_id, String request_final_usage_json_str);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(RequestStreamOutput, ObjectRef,
                                             RequestStreamOutputObj);
};

}  // namespace serve
}  // namespace llm
}  // namespace mlc

#endif  // MLC_LLM_SERVE_DATA_H_