#>cpp
/// @file

#pragma once

#include "query/frontend/ast/ast_serialization.hpp"
#include "query/plan/distributed_ops.capnp.h"
#include "query/plan/operator.hpp"
#include "query/serialization.hpp"
#include "storage/distributed/rpc/serialization.hpp"
cpp<#

(load "query/plan/operator.lcp")

(lcp:namespace query)
(lcp:namespace plan)

#>cpp
class PullRemote;
class Synchronize;
class PullRemoteOrderBy;
class DistributedExpand;
class DistributedExpandBfs;
class DistributedCreateNode;
class DistributedCreateExpand;

using DistributedOperatorCompositeVisitor =
    ::utils::CompositeVisitor<PullRemote, Synchronize, PullRemoteOrderBy,
                              DistributedExpand, DistributedExpandBfs,
                              DistributedCreateNode, DistributedCreateExpand>;

/// Base class for visiting regular and distributed LogicalOperator instances.
///
/// HierarchicalLogicalOperatorVisitor is inherited virtually, so that potential
/// multiple inheritance of DistributedOperatorVisitor and other
/// HierarchicalLogicalOperatorVisitor derived types is possible. Note that
/// virtual inheritance resolves the diamond problem, but this still carries a
/// cost. For example, you can no longer use static_cast to downcast a type even
/// though you are 100% sure downcast would be correct. dynamic_cast should work
/// as usual.
class DistributedOperatorVisitor : public virtual HierarchicalLogicalOperatorVisitor,
                                   public DistributedOperatorCompositeVisitor {
 public:
  using DistributedOperatorCompositeVisitor::PostVisit;
  using DistributedOperatorCompositeVisitor::PreVisit;
  using HierarchicalLogicalOperatorVisitor::PostVisit;
  using HierarchicalLogicalOperatorVisitor::PreVisit;
  using HierarchicalLogicalOperatorVisitor::Visit;
};
cpp<#

(lcp:define-class pull-remote (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (plan-id :int64_t :initval 0 :scope :public)
   (symbols "std::vector<Symbol>" :scope :public
            :capnp-save (lcp:capnp-save-vector "::query::capnp::Symbol" "Symbol")
            :capnp-load (lcp:capnp-load-vector "::query::capnp::Symbol" "Symbol")))
  (:documentation
   "An operator in distributed Memgraph that yields both local and remote (from
other workers) frames. Obtaining remote frames is done through RPC calls to
`distributed::ProduceRpcServer`s running on all the workers.

This operator aims to yield results as fast as possible and lose minimal
time on data transfer. It gives no guarantees on result order.")
  (:public
   #>cpp
   PullRemote() {}
   PullRemote(const std::shared_ptr<LogicalOperator> &input, int64_t plan_id,
              const std::vector<Symbol> &symbols)
       : input_(input), plan_id_(plan_id), symbols_(symbols) {}
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> OutputSymbols(const SymbolTable &) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(defun slk-load-pull-remote (member)
  #>cpp
  std::shared_ptr<query::plan::LogicalOperator> op;
  slk::Load<query::plan::LogicalOperator>(
      &op, reader, &helper->loaded_ops,
      [&helper](auto *op, auto *reader) {
        slk::ConstructAndLoad(op, reader, helper);
      });
  self->${member} = std::static_pointer_cast<query::plan::PullRemote>(op);
  cpp<#)

(defun load-pull-remote (reader member-name capnp-name)
  (declare (ignore capnp-name))
  #>cpp
  ${member-name} = std::static_pointer_cast<PullRemote>(
      utils::LoadSharedPtr<capnp::LogicalOperator, LogicalOperator>(${reader},
                       [helper](const auto &reader) {
                         std::unique_ptr<LogicalOperator> op;
                         Load(&op, reader, helper);
                         return op.release();
                       }, &helper->loaded_ops));
  cpp<#)

(lcp:define-class synchronize (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (pull-remote "std::shared_ptr<PullRemote>" :scope :public
                :slk-save #'slk-save-operator-pointer
                :slk-load #'slk-load-pull-remote
                :capnp-save #'save-operator-pointer
                :capnp-load #'load-pull-remote
                :clone (lambda (source dest)
                         #>cpp
                         if (${source}) {
                           std::shared_ptr<LogicalOperator> tmp = ${source}->Clone(storage);
                           ${dest} = std::static_pointer_cast<PullRemote>(tmp);
                         } else {
                           ${dest} = nullptr;
                         }
                         cpp<#))
   (advance-command :bool :initval "false" :scope :public))
  (:documentation
   "Operator used to synchronize stages of plan execution between the master and
all the workers. Synchronization is necessary in queries that update that
graph state because updates (as well as creations and deletions) are deferred
to avoid multithreaded modification of graph element data (as it's not
thread-safe).

Logic of the synchronize operator is:

1. If there is a Pull, tell all the workers to pull on that plan and
   accumulate results without sending them to the master. This is async.
2. Accumulate local results, in parallel with 1. getting executed on workers.
3. Wait till the master and all the workers are done accumulating.
4. Advance the command, if necessary.
5. Tell all the workers to apply their updates. This is async.
6. Apply local updates, in parallel with 5. on the workers.
7. Notify workers that the command has advanced, if necessary.
8. Yield all the results, first local, then from Pull if available.")
  (:public
   #>cpp
   Synchronize() {}
   Synchronize(const std::shared_ptr<LogicalOperator> &input,
               const std::shared_ptr<PullRemote> &pull_remote,
               bool advance_command)
       : input_(input),
         pull_remote_(pull_remote),
         advance_command_(advance_command) {}
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   std::vector<Symbol> OutputSymbols(
       const SymbolTable &symbol_table) const override {
     return input_->OutputSymbols(symbol_table);
   }

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:define-class pull-remote-order-by (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (plan-id :int64_t :initval 0 :scope :public)
   (symbols "std::vector<Symbol>" :scope :public
            :capnp-save (lcp:capnp-save-vector "::query::capnp::Symbol" "Symbol")
            :capnp-load (lcp:capnp-load-vector "::query::capnp::Symbol" "Symbol"))
   (order-by "std::vector<Expression *>" :scope :public
             :slk-save #'slk-save-ast-vector
             :slk-load (slk-load-ast-vector "Expression")
             :capnp-type "List(Ast.Tree)"
             :capnp-save (save-ast-vector "Expression *")
             :capnp-load (load-ast-vector "Expression *"))
   (compare "TypedValueVectorCompare" :scope :public
            :capnp-type "Query.TypedValueVectorCompare"))
  (:documentation
   "Operator that merges distributed OrderBy operators.
Instead of using a regular OrderBy on master (which would collect all remote
results and order them), we can have each worker do an OrderBy locally and
have the master rely on the fact that the results are ordered and merge them
by having only one result from each worker.")
  (:public
   #>cpp
   PullRemoteOrderBy() {}
   PullRemoteOrderBy(
       const std::shared_ptr<LogicalOperator> &input, int64_t plan_id,
       const std::vector<SortItem> &order_by, const std::vector<Symbol> &symbols);
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;

   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;
   std::vector<Symbol> OutputSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:define-class distributed-expand (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (input-symbol "Symbol" :scope :public)
   (common "ExpandCommon" :scope :public)
   (graph-view "GraphView" :scope :public
               :capnp-init nil
               :capnp-save (lcp:capnp-save-enum "::query::capnp::GraphView"
                                                "GraphView"
                                                '(old new))
               :capnp-load (lcp:capnp-load-enum "::query::capnp::GraphView"
                                                "GraphView"
                                                '(old new))
               :documentation
               "State from which the input node should get expanded."))
  (:documentation "Distributed version of Expand operator")
  (:public
   #>cpp
   DistributedExpand() {}
   DistributedExpand(const std::shared_ptr<LogicalOperator> &input,
                     Symbol input_symbol, Symbol node_symbol, Symbol edge_symbol,
                     EdgeAtom::Direction direction,
                     const std::vector<storage::EdgeType> &edge_types,
                     bool existing_node, GraphView graph_view);
   DistributedExpand(const std::shared_ptr<LogicalOperator> &input,
                     Symbol input_symbol, const ExpandCommon &common);

   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:define-class distributed-expand-bfs (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (input-symbol "Symbol" :scope :public)
   (common "ExpandCommon" :scope :public)
   (lower-bound "Expression *" :scope :public
                :documentation "Optional lower bound, default is 1"
                :slk-save #'slk-save-ast-pointer
                :slk-load (slk-load-ast-pointer "Expression")
                :capnp-type "Ast.Tree" :capnp-init nil
                :capnp-save #'save-ast-pointer
                :capnp-load (load-ast-pointer "Expression *"))
   (upper-bound "Expression *" :scope :public
                :documentation "Optional upper bound, default is infinity"
                :slk-save #'slk-save-ast-pointer
                :slk-load (slk-load-ast-pointer "Expression")
                :capnp-type "Ast.Tree" :capnp-init nil
                :capnp-save #'save-ast-pointer
                :capnp-load (load-ast-pointer "Expression *"))
   (filter-lambda "ExpansionLambda" :scope :public
                  :documentation "Filter that must be satisfied for expansion to succeed."
                  :slk-load (lambda (member)
                              #>cpp
                              slk::Load(&self->${member}, reader, &helper->ast_storage);
                              cpp<#)
                  :capnp-type "ExpansionLambda"
                  :capnp-load (lambda (reader member capnp-name)
                                #>cpp
                                Load(&${member}, ${reader}, &helper->ast_storage);
                                cpp<#)))
  (:documentation "BFS expansion operator suited for distributed execution.")
  (:public
   #>cpp
   DistributedExpandBfs() {}
   DistributedExpandBfs(const std::shared_ptr<LogicalOperator> &input,
                        Symbol input_symbol, Symbol node_symbol,
                        Symbol edge_symbol, EdgeAtom::Direction direction,
                        const std::vector<storage::EdgeType> &edge_types,
                        bool existing_node, Expression *lower_bound,
                        Expression *upper_bound,
                        const ExpansionLambda &filter_lambda);
   DistributedExpandBfs(const std::shared_ptr<LogicalOperator> &input,
                        Symbol input_symbol, const ExpandCommon &common,
                        Expression *lower_bound, Expression *upper_bound,
                        const ExpansionLambda &filter_lambda);

   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:define-class distributed-create-node (logical-operator)
  ((input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (node-info "NodeCreationInfo" :scope :public
              :slk-save (lambda (m)
                          #>cpp
                          slk::Save(self.${m}, builder, helper);
                          cpp<#)
              :slk-load (lambda (m)
                          #>cpp
                          slk::Load(&self->${m}, reader, helper);
                          cpp<#))
   (on-random-worker :bool :initval "false" :scope :public))
  (:documentation "Create nodes in distributed environment.")
  (:public
   #>cpp
   DistributedCreateNode() {}
   DistributedCreateNode(const std::shared_ptr<LogicalOperator> &input,
                         const NodeCreationInfo &node_info, bool on_random_worker);

   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:define-class distributed-create-expand (logical-operator)
  ((node-info "NodeCreationInfo" :scope :public
              :slk-save (lambda (m)
                          #>cpp
                          slk::Save(self.${m}, builder, helper);
                          cpp<#)
              :slk-load (lambda (m)
                          #>cpp
                          slk::Load(&self->${m}, reader, helper);
                          cpp<#))
   (edge-info "EdgeCreationInfo" :scope :public
              :slk-save (lambda (m)
                          #>cpp
                          slk::Save(self.${m}, builder, helper);
                          cpp<#)
              :slk-load (lambda (m)
                          #>cpp
                          slk::Load(&self->${m}, reader, helper);
                          cpp<#))
   (input "std::shared_ptr<LogicalOperator>" :scope :public
          :slk-save #'slk-save-operator-pointer
          :slk-load #'slk-load-operator-pointer
          :capnp-save #'save-operator-pointer
          :capnp-load #'load-operator-pointer)
   (input-symbol "Symbol" :scope :public)
   (existing-node :bool :scope :public))
  (:documentation "Distributed version of CreateExpand")
  (:public
   #>cpp
   DistributedCreateExpand() {}
   DistributedCreateExpand(const NodeCreationInfo &node_info,
                           const EdgeCreationInfo &edge_info,
                           const std::shared_ptr<LogicalOperator> &input,
                           Symbol input_symbol, bool existing_node);
   bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
   std::unique_ptr<Cursor> MakeCursor(
       database::GraphDbAccessor &db) const override;
   std::vector<Symbol> ModifiedSymbols(const SymbolTable &) const override;

   bool HasSingleInput() const override { return true; }
   std::shared_ptr<LogicalOperator> input() const override { return input_; }
   void set_input(std::shared_ptr<LogicalOperator> input) override {
     input_ = input;
   }
   cpp<#)
  (:serialize (:slk) (:capnp))
  (:clone))

(lcp:pop-namespace) ;; plan
(lcp:pop-namespace) ;; query
