#>cpp
#pragma once

#include <cstdint>
#include <functional>
#include <string>

#include "communication/rpc/messages.hpp"
#include "distributed/pull_produce_rpc_messages.capnp.h"
#include "query/context.hpp"
#include "query/frontend/semantic/symbol.hpp"
#include "query/parameters.hpp"
#include "query/serialization.hpp"
#include "storage/distributed/address_types.hpp"
#include "transactions/type.hpp"
#include "rpc/serialization.hpp"
cpp<#

(lcp:in-impl
 #>cpp
 #include "database/distributed/graph_db.hpp"
 #include "database/distributed/graph_db_accessor.hpp"
 #include "distributed/data_manager.hpp"
 cpp<#)

(load "transactions/distributed/serialization.lcp")

(lcp:namespace distributed)

(lcp:capnp-namespace "distributed")

(lcp:capnp-import 'storage "/storage/distributed/rpc/serialization.capnp")
(lcp:capnp-import 'query "/query/serialization.capnp")
(lcp:capnp-import 'sem "/query/frontend/semantic/symbol_serialization.capnp")
(lcp:capnp-import 'utils "/rpc/serialization.capnp")

(lcp:capnp-type-conversion "tx::CommandId" "UInt32")
(lcp:capnp-type-conversion "tx::TransactionId" "UInt64")

#>cpp
// Forward declare for LoadGraphElement.
class DataManager;

/// The default number of results returned via RPC from remote execution to the
/// master that requested it.
constexpr int kDefaultBatchSize = 20;
cpp<#

(lcp:define-enum pull-state
    (cursor-exhausted
     cursor-in-progress
     serialization-error
     lock-timeout-error
     update-deleted-error
     reconstruction-error
     unable-to-delete-vertex-error
     hinted-abort-error
     query-error)
  (:documentation "Returned along with a batch of results in the remote-pull
RPC. Indicates the state of execution on the worker.")
  (:serialize))

(lcp:define-struct pull-data ()
  ((pull-state "PullState")
   (frames "std::vector<std::vector<query::TypedValue>>"))
  (:documentation
   "The data returned to the end consumer (the Pull operator). Contains only
the relevant parts of the response, ready for use."))

(defun slk-save-frames (member)
  #>cpp
  size_t frame_count = self.${member}.size();
  slk::Save(frame_count, builder);
  for (const auto &frame : self.${member}) {
    size_t frame_size = frame.size();
    slk::Save(frame_size, builder);
    for (const auto &value : frame) {
      slk::Save(value, builder, self.send_versions, self.worker_id);
    }
  }
  cpp<#)

(defun slk-load-frames (member)
  #>cpp
  size_t frame_count = 0;
  slk::Load(&frame_count, reader);
  self->${member}.reserve(frame_count);
  for (size_t frame_i = 0; frame_i < frame_count; ++frame_i) {
    size_t frame_size = 0;
    slk::Load(&frame_size, reader);
    std::vector<query::TypedValue> frame(frame_size);
    for (size_t val_i = 0; val_i < frame_size; ++val_i) {
      slk::Load(&frame[val_i], reader, dba, data_manager);
    }
    self->${member}.emplace_back(std::move(frame));
  }
  cpp<#)

(lcp:define-struct pull-res-data ()
  ((pull-state "PullState")
   (frames "std::vector<std::vector<query::TypedValue>>"
           :slk-save #'slk-save-frames
           :slk-load #'slk-load-frames
           :capnp-type "List(List(Query.TypedValue))"
           :capnp-save
           (lambda (builder member capnp-name)
             (declare (ignore capnp-name))
             #>cpp
             for (size_t frame_i = 0; frame_i < ${member}.size(); ++frame_i) {
               const auto &frame = ${member}[frame_i];
               auto frame_builder = ${builder}.init(frame_i, frame.size());
               for (size_t val_i = 0; val_i < frame.size(); ++val_i) {
                 const auto &value = frame[val_i];
                 auto value_builder = frame_builder[val_i];
                 query::SaveCapnpTypedValue(value, &value_builder, self.send_versions, self.worker_id);
               }
             }
             cpp<#)
           :capnp-load
           (lambda (reader member capnp-name)
             (declare (ignore capnp-name))
             #>cpp
             ${member}.reserve(${reader}.size());
             for (const auto &frame_reader : ${reader}) {
               std::vector<query::TypedValue> current_frame;
               current_frame.reserve(frame_reader.size());
               for (const auto &value_reader : frame_reader) {
                 query::TypedValue value;
                 query::LoadCapnpTypedValue(value_reader, &value, dba, data_manager);
                 current_frame.emplace_back(value);
               }
               ${member}.emplace_back(current_frame);
             }
             cpp<#))
   (worker-id :int16_t :dont-save t
              :documentation
              "Id of the worker on which the response is created, used for
serializing vertices (converting local to global addresses). Indicates which
of (old, new) records of a graph element should be sent.")
   (send-versions "storage::SendVersions" :dont-save t)
   ;; Temporary caches used between deserialization and post-processing
   ;; (transfering the ownership of this data to a Cache).
   (vertices "std::vector<GraphElementData<Vertex>>" :dont-save t)
   (edges "std::vector<GraphElementData<Edge>>" :dont-save t)
   (paths "std::vector<PathData>" :dont-save t))
  (:documentation
   "The data of the remote pull response. Post-processing is required after
deserialization to initialize Vertex/Edge typed values in the frames (possibly
encapsulated in lists/maps) to their proper values. This requires a
GraphDbAccessor and therefore can't be done as part of deserialization.

TODO - make it possible to inject a &GraphDbAcessor from the Pull layer all
the way into RPC data deserialization to remove the requirement for
post-processing. The current approach of holding references to parts of the
frame (potentially embedded in lists/maps) is too error-prone.")
  (:public
   #>cpp
   private:
   cpp<#
   (lcp:define-struct (graph-element-data t-record) ()
     ((cypher-id :int64_t)
      (global-address "storage::Address<mvcc::VersionList<TRecord>>")
      (old-record "std::unique_ptr<TRecord>")
      (new-record "std::unique_ptr<TRecord>")
      (element-in-frame
       "query::TypedValue *"
       :documentation
       "The position in frame is optional. This same structure is used for
deserializing path elements, in which case the vertex/edge in question is not
directly part of the frame."))
     (:documentation
      "Temp cache for deserialized vertices and edges. These objects are
created during deserialization. They are used immediatelly after during
post-processing. The vertex/edge data ownership gets transfered to the Cache,
and the `element_in_frame` reference is used to set the appropriate accessor
to the appropriate value. Not used on side that generates the response.")
     (:public
      #>cpp
      GraphElementData(int64_t cypher_id, storage::Address<mvcc::VersionList<TRecord>> address,
                       std::unique_ptr<TRecord> old_record, std::unique_ptr<TRecord> new_record,
                       query::TypedValue *element_in_frame)
        : cypher_id(cypher_id),
          global_address(address),
          old_record(std::move(old_record)),
          new_record(std::move(new_record)),
          element_in_frame(element_in_frame) {}
      cpp<#))
   (lcp:define-struct path-data ()
     ((vertices "std::vector<GraphElementData<Vertex>>")
      (edges "std::vector<GraphElementData<Edge>>")
      (path-in-frame "query::TypedValue *"))
     (:public
      #>cpp
      PathData(query::TypedValue *path_in_frame) : path_in_frame(path_in_frame) {}
      cpp<#)
     (:documentation "Same like `GraphElementData`, but for paths."))
   #>cpp
   public:
    PullResData() {}  // Default constructor required for serialization.
    PullResData(int worker_id, storage::SendVersions send_versions)
        : worker_id(worker_id), send_versions(send_versions) {}

    PullResData(const PullResData &) = delete;
    PullResData &operator=(const PullResData &) = delete;
    PullResData(PullResData &&) = default;
    PullResData &operator=(PullResData &&) = default;
   cpp<#)
  (:serialize (:slk :load-args '((dba "database::GraphDbAccessor *")
                                 (data-manager "distributed::DataManager *")))
              (:capnp :load-args '((dba "database::GraphDbAccessor *")
                                   (data-manager "distributed::DataManager *")))))

(lcp:define-rpc pull
    (:request
     ((tx-id "tx::TransactionId")
      (tx-snapshot "tx::Snapshot"
                   :slk-save #'slk-save-snapshot
                   :slk-load #'slk-load-snapshot
                   :capnp-type "List(UInt64)"
                   :capnp-init nil
                   :capnp-save #'save-snapshot
                   :capnp-load #'load-snapshot)
      (plan-id :int64_t)
      (command-id "tx::CommandId")
      (timestamp :int64_t)
      (parameters "query::Parameters"
                  :capnp-type "Utils.Map(Utils.BoxInt64, Storage.PropertyValue)")
      (symbols "std::vector<query::Symbol>" :capnp-type "List(Sem.Symbol)")
      (accumulate :bool)
      (batch-size :int64_t)
      ;; Indicates which of (old, new) records of a graph element should be sent.
      (send-versions "storage::SendVersions"
                     :capnp-type "Storage.SendVersions" :capnp-init nil
                     :capnp-save (lcp:capnp-save-enum "storage::capnp::SendVersions"
                                                      "storage::SendVersions"
                                                      '(both only-old only-new))
                     :capnp-load (lcp:capnp-load-enum "storage::capnp::SendVersions"
                                                      "storage::SendVersions"
                                                      '(both only-old only-new)))))
  (:response
   ((data "PullResData" :initarg :move
          :slk-load (lambda (m)
                      #>cpp
                      slk::Load(&self->${m}, reader, dba, data_manager);
                      cpp<#)))
   (:serialize (:slk :load-args '((dba "database::GraphDbAccessor *")
                                  (data-manager "distributed::DataManager *")))
               (:capnp :load-args '((dba "database::GraphDbAccessor *")
                                    (data-manager "distributed::DataManager *"))))))

;; TODO make a separate RPC for the continuation of an existing pull, as an
;; optimization not to have to send the full PullReqData pack every time.

(lcp:define-rpc reset-cursor
    (:request
     ((tx-id "tx::TransactionId")
      (plan-id :int64_t)
      (command-id "tx::CommandId")))
    (:response ()))

(lcp:define-rpc transaction-command-advanced
    (:request ((member "tx::TransactionId")))
  (:response ()))

(lcp:pop-namespace) ;; distributed
