#include "module.h"

#include <ctranslate2/encoder.h>

#include "replica_pool.h"

namespace ctranslate2 {
  namespace python {

    class EncoderWrapper : public ReplicaPoolHelper<Encoder> {
    public:
      using ReplicaPoolHelper::ReplicaPoolHelper;

      EncoderForwardOutput
      forward_batch(const std::variant<BatchTokens, BatchIds, StorageView>& inputs,
                    const std::optional<StorageView>& lengths,
                    const std::optional<BatchIds>& token_type_ids) {
        std::future<EncoderForwardOutput> future;
        std::shared_lock lock(_mutex);
        assert_model_is_ready();

        switch (inputs.index()) {
        case 0:
          future = _pool->forward_batch_async(
            std::get<BatchTokens>(inputs),
            token_type_ids.value_or(std::vector<std::vector<size_t>>()));
          break;
        case 1:
          future = _pool->forward_batch_async(
            std::get<BatchIds>(inputs),
            token_type_ids.value_or(std::vector<std::vector<size_t>>()));
          break;
        case 2:
          if (!lengths)
            throw std::invalid_argument("lengths vector is required when passing a dense input");
          future = _pool->forward_batch_async(
            std::get<StorageView>(inputs),
            lengths.value(),
            token_type_ids.value_or(std::vector<std::vector<size_t>>()));
          break;
        }

        return future.get();
      }
    };


    void register_encoder(py::module& m) {
      py::class_<EncoderForwardOutput>(m, "EncoderForwardOutput",
                                       "Forward output of an encoder model.")

        .def_readonly("last_hidden_state", &EncoderForwardOutput::last_hidden_state,
                      "Output of the last layer.")
        .def_readonly("pooler_output", &EncoderForwardOutput::pooler_output,
                      "Output of the pooling layer.")

        .def("__repr__", [](const EncoderForwardOutput& output) {
          return "EncoderForwardOutput(last_hidden_state="
            + std::string(py::repr(py::cast(output.last_hidden_state)))
            + ", pooler_output=" + std::string(py::repr(py::cast(output.pooler_output)))
            + ")";
        })
        ;

      py::class_<EncoderWrapper>(
        m, "Encoder",
        R"pbdoc(
            A text encoder.

            Example:

                >>> encoder = ctranslate2.Encoder("model/", device="cpu")
                >>> encoder.forward_batch([["▁Hello", "▁world", "!"]])
        )pbdoc")

        .def(py::init<const std::string&, const std::string&, const std::variant<int, std::vector<int>>&, const StringOrMap&, size_t, size_t, long, bool, bool, py::object>(),
             py::arg("model_path"),
             py::arg("device")="cpu",
             py::kw_only(),
             py::arg("device_index")=0,
             py::arg("compute_type")="default",
             py::arg("inter_threads")=1,
             py::arg("intra_threads")=0,
             py::arg("max_queued_batches")=0,
             py::arg("flash_attention")=false,
             py::arg("tensor_parallel")=false,
             py::arg("files")=py::none(),
             R"pbdoc(
                 Initializes the encoder.

                 Arguments:
                   model_path: Path to the CTranslate2 model directory.
                   device: Device to use (possible values are: cpu, cuda, auto).
                   device_index: Device IDs where to place this encoder on.
                   compute_type: Model computation type or a dictionary mapping a device name
                     to the computation type (possible values are: default, auto, int8, int8_float32,
                     int8_float16, int8_bfloat16, int16, float16, bfloat16, float32).
                   inter_threads: Maximum number of parallel generations.
                   intra_threads: Number of OpenMP threads per encoder (0 to use a default value).
                   max_queued_batches: Maximum numbers of batches in the queue (-1 for unlimited,
                     0 for an automatic value). When the queue is full, future requests will block
                     until a free slot is available.
                   flash_attention: run model with flash attention 2 for self-attention layer
                   tensor_parallel: run model with tensor parallel mode
                   files: Load model files from the memory. This argument is a dictionary mapping
                     file names to file contents as file-like or bytes objects. If this is set,
                     :obj:`model_path` acts as an identifier for this model.
             )pbdoc")

        .def_property_readonly("device", &EncoderWrapper::device,
                               "Device this encoder is running on.")
        .def_property_readonly("device_index", &EncoderWrapper::device_index,
                               "List of device IDs where this encoder is running on.")
        .def_property_readonly("compute_type", &EncoderWrapper::compute_type,
                               "Computation type used by the model.")
        .def_property_readonly("num_encoders", &EncoderWrapper::num_replicas,
                               "Number of encoders backing this instance.")
        .def_property_readonly("num_queued_batches", &EncoderWrapper::num_queued_batches,
                               "Number of batches waiting to be processed.")
        .def_property_readonly("tensor_parallel", &EncoderWrapper::tensor_parallel,
                               "Run model with tensor parallel mode.")
        .def_property_readonly("num_active_batches", &EncoderWrapper::num_active_batches,
                               "Number of batches waiting to be processed or currently processed.")

        .def("forward_batch", &EncoderWrapper::forward_batch,
             py::arg("inputs"),
             py::arg("lengths")=py::none(),
             py::arg("token_type_ids")=py::none(),
             py::call_guard<py::gil_scoped_release>(),
             R"pbdoc(
                 Forwards a batch of sequences in the encoder.

                 Arguments:
                   inputs: A batch of sequences either as string tokens or token IDs.
                     This argument can also be a dense int32 array with shape
                     ``[batch_size, max_length]`` (e.g. created from a Numpy array or PyTorch tensor).
                   lengths: The length of each sequence as a int32 array with shape
                     ``[batch_size]``. Required when :obj:`inputs` is a dense array.
                   token_type_ids: A batch of token type IDs of same shape as :obj:`inputs`.
                     ``[batch_size, max_length]``.

                 Returns:
                   The encoder model output.
             )pbdoc")

        .def("unload_model", &EncoderWrapper::unload_model,
             py::arg("to_cpu")=false,
             py::call_guard<py::gil_scoped_release>(),
             R"pbdoc(
                 Unloads the model attached to this encoder but keep enough runtime context
                 to quickly resume encoder on the initial device.

                 Arguments:
                   to_cpu: If ``True``, the model is moved to the CPU memory and not fully unloaded.
             )pbdoc")

        .def("load_model", &EncoderWrapper::load_model,
             py::arg("keep_cache")=false,
             py::call_guard<py::gil_scoped_release>(),
             R"pbdoc(
                 Loads the model back to the initial device.

                 Arguments:
                   keep_cache: If ``True``, the model cache in the CPU memory is not deleted if it exists.
             )pbdoc")

        .def_property_readonly("model_is_loaded", &EncoderWrapper::model_is_loaded,
                               "Whether the model is loaded on the initial device and ready to be used.")
        ;
    }

  }
}