text: /**
   *
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  #ifndef EXTENSIONS_JNI_JVM_REFERNCEOBJECTS_H_
  #define EXTENSIONS_JNI_JVM_REFERNCEOBJECTS_H_
  
  #include <string>
  #include <vector>
  #include <sstream>
  #include <iterator>
  #include <algorithm>
  #include <jni.h>
  #include "JavaServicer.h"
  #include "core/Processor.h"
  #include "core/ProcessSession.h"
  #include "core/WeakReference.h"
  #include "utils/gsl.h"
  
  namespace org {
  namespace apache {
  namespace nifi {
  namespace minifi {
  namespace jni {
  
  /**
   * Represents a flow file. Exists to provide the ability to remove
   * global references amongst the
   */
  class JniFlowFile : public core::WeakReference {
   public:
    JniFlowFile(std::shared_ptr<core::FlowFile> ref, const std::shared_ptr<JavaServicer> &servicer, jobject ff)
        : removed(false),
          ff_object(ff),
          ref_(ref),
          servicer_(servicer) {
    }
  
    virtual ~JniFlowFile() = default;
  
    void remove() override;
  
    std::shared_ptr<core::FlowFile> get() const {
      return ref_;
    }
  
    jobject getJniReference() {
      return ff_object;
    }
  
    bool operator==(const JniFlowFile &other) const {
      // compare the pointers
      return ref_ == other.ref_;
    }
  
    bool empty() const {
      return removed;
    }
  
   protected:
  
    bool removed;
  
    jobject ff_object;
  
    std::mutex session_mutex_;
  
    std::shared_ptr<core::FlowFile> ref_;
  
    std::shared_ptr<JavaServicer> servicer_;
  };
  
  /**
   * Quick check to determine if a FF is empty.
   */
  struct check_empty_ff : public std::unary_function<std::shared_ptr<JniFlowFile>, bool> {
    bool operator()(std::shared_ptr<JniFlowFile> session) const {
      return session->empty();
    }
  };
  
  class JniByteOutStream : public minifi::OutputStreamCallback {
   public:
    JniByteOutStream(jbyte *bytes, size_t length)
        : bytes_(bytes),
          length_(length) {
    }
  
    virtual ~JniByteOutStream() = default;
    virtual int64_t process(const std::shared_ptr<minifi::io::BaseStream>& stream) {
      const auto write_ret = stream->write((uint8_t*)bytes_, length_);
      return io::isError(write_ret) ? -1 : gsl::narrow<int64_t>(write_ret);
    }
   private:
    jbyte *bytes_;
    size_t length_;
  };
  
  /**
   * Jni byte input stream
   */
  class JniByteInputStream : public minifi::InputStreamCallback {
   public:
    JniByteInputStream(uint64_t size)
        : stream_(nullptr),
          read_size_(0) {
      buffer_size_ = size;
      buffer_ = new uint8_t[buffer_size_];
    }
    ~JniByteInputStream() {
      if (buffer_)
        delete[] buffer_;
    }
    int64_t process(const std::shared_ptr<minifi::io::BaseStream>& stream) {
      stream_ = stream;
      return 0;
    }
  
    int64_t read(JNIEnv *env, jbyteArray arr, int offset, int size) {
      gsl_Expects(size >= 0);
      if (stream_ == nullptr) {
        return -1;
      }
  
      // seek to offset
      uint64_t remaining = gsl::narrow<uint64_t>(size);
      int writtenOffset = 0;
      int read = 0;
      do {
        // JNI takes size as int, there's not much we can do here to support 2GB+ sizes
        int actual = static_cast<int>(stream_->read(buffer_, std::min(remaining, buffer_size_)));
        if (actual <= 0) {
          if (read == 0) {
            stream_ = nullptr;
            return -1;
          }
          break;
        }
  
        read += actual;
        env->SetByteArrayRegion(arr, offset + writtenOffset, actual, (jbyte*) buffer_);
        writtenOffset += actual;
  
        remaining -= actual;
      } while (remaining > 0);
  
      return read;
    }
  
    int64_t read(char &arr) {
      return stream_->read(arr);
    }
  
    std::shared_ptr<minifi::io::BaseStream> stream_;
    uint8_t *buffer_;
    uint64_t buffer_size_;
    uint64_t read_size_;
  };
  
  class JniInputStream : public core::WeakReference {
   public:
    JniInputStream(std::unique_ptr<JniByteInputStream> jbi, jobject in_instance, const std::shared_ptr<JavaServicer> &servicer)
        : removed_(false),
          in_instance_(in_instance),
          jbi_(std::move(jbi)),
          servicer_(servicer) {
    }
  
    void remove() override {
      std::lock_guard<std::mutex> guard(mutex_);
      if (!removed_) {
        servicer_->attach()->DeleteGlobalRef(in_instance_);
        removed_ = true;
        jbi_ = nullptr;
      }
    }
  
    int64_t read(char &arr) {
      if (!removed_) {
        return jbi_->read(arr);
      }
      return -1;
    }
  
    int64_t read(JNIEnv *env, jbyteArray arr, int offset, int size) {
      if (!removed_) {
        return jbi_->read(env, arr, offset, size);
      }
      return -1;
    }
  
   private:
    std::mutex mutex_;
    bool removed_;
    jobject in_instance_;
    std::unique_ptr<JniByteInputStream> jbi_;
    std::shared_ptr<JavaServicer> servicer_;
  };
  
  class JniSession : public core::WeakReference {
   public:
    JniSession(const std::shared_ptr<core::ProcessSession> &session, jobject session_instance, const std::shared_ptr<JavaServicer> &servicer)
        : removed_(false),
          session_instance_(session_instance),
          session_(session),
          servicer_(servicer) {
    }
  
    void remove() override {
      std::lock_guard<std::mutex> guard(session_mutex_);
      if (!removed_) {
        for (auto ff : global_ff_objects_) {
          ff->remove();
        }
  
        for (auto in : input_streams_) {
          in->remove();
        }
        global_ff_objects_.clear();
        servicer_->attach()->DeleteGlobalRef(session_instance_);
        removed_ = true;
      }
    }
  
    std::shared_ptr<core::ProcessSession> &getSession() {
      return session_;
    }
  
    JniFlowFile *getFlowFileReference(const std::shared_ptr<core::FlowFile> &ff) {
      for (auto &jni_ff : global_ff_objects_) {
        if (jni_ff->get().get() == ff.get()) {
          return jni_ff.get();
        }
      }
      return nullptr;
    }
  
    JniFlowFile * addFlowFile(std::shared_ptr<JniFlowFile> ff) {
      std::lock_guard<std::mutex> guard(session_mutex_);
      global_ff_objects_.push_back(ff);
      return ff.get();
    }
  
    void addInputStream(std::shared_ptr<JniInputStream> in) {
      std::lock_guard<std::mutex> guard(session_mutex_);
      input_streams_.push_back(in);
    }
  
    std::shared_ptr<JavaServicer> getServicer() const {
      return servicer_;
    }
  
    bool prune() {
      global_ff_objects_.erase(std::remove_if(global_ff_objects_.begin(), global_ff_objects_.end(), check_empty_ff()), global_ff_objects_.end());
      if (global_ff_objects_.empty()) {
        remove();
      }
      return global_ff_objects_.empty();
    }
    bool empty() const {
      std::lock_guard<std::mutex> guard(session_mutex_);
      for (auto ff : global_ff_objects_) {
        if (!ff->empty())
          return false;
      }
      return true;
    }
  
   protected:
    bool removed_;
    mutable std::mutex session_mutex_;
    jobject session_instance_;
    std::shared_ptr<core::ProcessSession> session_;
    std::shared_ptr<JavaServicer> servicer_;
    std::vector<std::shared_ptr<JniInputStream>> input_streams_;
    // we own
    std::vector<std::shared_ptr<JniFlowFile>> global_ff_objects_;
  };
  
  struct check_empty : public std::unary_function<std::shared_ptr<JniSession>, bool> {
    bool operator()(std::shared_ptr<JniSession> session) const {
      return session->prune();
    }
  };
  
  class JniSessionFactory : public core::WeakReference {
   public:
  
    JniSessionFactory(const std::shared_ptr<core::ProcessSessionFactory> &factory, const std::shared_ptr<JavaServicer> &servicer, jobject java_object)
        : servicer_(servicer),
          factory_(factory),
          java_object_(java_object) {
    }
  
    void remove() override {
      std::lock_guard<std::mutex> guard(session_mutex_);
      // remove all of the sessions
      // this should spark their destructor
      for (auto session : sessions_) {
        session->remove();
      }
      sessions_.clear();
  
      if (java_object_) {
        // detach the global reference and let Java take care of this object.
        servicer_->attach()->DeleteGlobalRef(java_object_);
      }
    }
  
    jobject getJavaReference() const {
      return java_object_;
    }
  
    std::shared_ptr<JavaServicer> getServicer() const {
      return servicer_;
    }
  
    std::shared_ptr<core::ProcessSessionFactory> getFactory() const {
      return factory_;
    }
  
    /**
     */
    JniSession *addSession(std::shared_ptr<JniSession> session) {
      std::lock_guard<std::mutex> guard(session_mutex_);
  
      sessions_.erase(std::remove_if(sessions_.begin(), sessions_.end(), check_empty()), sessions_.end());
  
      sessions_.push_back(session);
  
      return session.get();
    }
  
   protected:
    std::shared_ptr<JavaServicer> servicer_;
    std::mutex session_mutex_;
    // we do not own this shared ptr
    std::shared_ptr<core::ProcessSessionFactory> factory_;
    // we own the sessions
    std::vector<std::shared_ptr<JniSession>> sessions_;
    // we own the java object
    jobject java_object_;
  };
  
  
  
  } /* namespace jni */
  } /* namespace minifi */
  } /* namespace nifi */
  } /* namespace apache */
  } /* namespace org */
  
  #endif /* EXTENSIONS_JNI_JVM_REFERNCEOBJECTS_H_ */
  