// 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.

#pragma once

#include <butil/macros.h>

#include <jni.h>

#include <limits.h>

#include <stddef.h>

#include <stdint.h>

#include <string>

#include <unordered_map>

#include "common/status.h"

#include "jni_md.h"

#include "util/defer_op.h"

#include "util/thrift_util.h"

#ifdef USE_HADOOP_HDFS

// defined in hadoop/hadoop-hdfs-project/hadoop-hdfs/src/main/native/libhdfs/jni_helper.c

extern "C" JNIEnv* getJNIEnv(void);

#endif

namespace doris {

class JniUtil;

#define RETURN_ERROR_IF_EXC(env)                     \

    do {                                             \

        if (env->ExceptionCheck()) [[unlikely]]      \

            return JniUtil::GetJniExceptionMsg(env); \

    } while (false)

#define JNI_CALL_METHOD_CHECK_EXCEPTION_DELETE_REF(type, result, env, func) \

    type result = env->func;                                                \

    DEFER(env->DeleteLocalRef(result));                                     \

    RETURN_ERROR_IF_EXC(env)

#define JNI_CALL_METHOD_CHECK_EXCEPTION(type, result, env, func) \

    type result = env->func;                                     \

    RETURN_ERROR_IF_EXC(env)

//In order to reduce the potential risks caused by not handling exceptions,

// you need to refer to  https://docs.oracle.com/javase/8/docs/technotes/guides/jni/spec/functions.html

// to confirm whether the jni method will throw an exception.

class JniUtil {

public:

    static Status Init() WARN_UNUSED_RESULT;

    static jmethodID throwable_to_string_id() { return throwable_to_string_id_; }

    static Status GetJNIEnv(JNIEnv** env) {

        if (tls_env_) {

            *env = tls_env_;

        } else {

            Status status = GetJNIEnvSlowPath(env);

            if (!status.ok()) {

                return status;

            }

        }

        if (*env == nullptr) {

            return Status::RuntimeError("Failed to get JNIEnv: it is nullptr.");

        }

        return Status::OK();

    }

    //jclass is generally a local reference.

    //Method ID and field ID values are forever.

    //If you want to use the jclass across multiple threads or multiple calls into the JNI code you need

    // to create a global reference to it with GetGlobalClassRef().

    static Status GetGlobalClassRef(JNIEnv* env, const char* class_str,

                                    jclass* class_ref) WARN_UNUSED_RESULT;

    static Status LocalToGlobalRef(JNIEnv* env, jobject local_ref,

                                   jobject* global_ref) WARN_UNUSED_RESULT;

    static Status GetJniExceptionMsg(JNIEnv* env, bool log_stack = true,

                                     const std::string& prefix = "") WARN_UNUSED_RESULT;

    static jclass jni_util_class() { return jni_util_cl_; }

    static jmethodID throwable_to_stack_trace_id() { return throwable_to_stack_trace_id_; }

    static const int64_t INITIAL_RESERVED_BUFFER_SIZE = 1024;

    // TODO: we need a heuristic strategy to increase buffer size for variable-size output.

    static inline int64_t IncreaseReservedBufferSize(int n) {

        return INITIAL_RESERVED_BUFFER_SIZE << n;

    }

    static Status get_jni_scanner_class(JNIEnv* env, const char* classname, jclass* loaded_class);

    static jobject convert_to_java_map(JNIEnv* env, const std::map<std::string, std::string>& map);

    static std::map<std::string, std::string> convert_to_cpp_map(JNIEnv* env, jobject map);

    static size_t get_max_jni_heap_memory_size();

    static Status clean_udf_class_load_cache(const std::string& function_signature);

private:

    static void parse_max_heap_memory_size_from_jvm(JNIEnv* env);

    static Status GetJNIEnvSlowPath(JNIEnv** env);

    static Status init_jni_scanner_loader(JNIEnv* env);

    static bool jvm_inited_;

    static jclass internal_exc_cl_;

    static jclass jni_native_method_exc_cl_;

    static jclass jni_util_cl_;

    static jmethodID throwable_to_string_id_;

    static jmethodID throwable_to_stack_trace_id_;

    static jmethodID get_jvm_metrics_id_;

    static jmethodID get_jvm_threads_id_;

    static jmethodID get_jmx_json_;

    // JNI scanner loader

    static jobject jni_scanner_loader_obj_;

    static jmethodID jni_scanner_loader_method_;

    // Thread-local cache of the JNIEnv for this thread.

    static __thread JNIEnv* tls_env_;

    static jlong max_jvm_heap_memory_size_;

    static jmethodID _clean_udf_cache_method_id;

};

/// Helper class for lifetime management of chars from JNI, releasing JNI chars when

/// destructed

class JniUtfCharGuard {

public:

    /// Construct a JniUtfCharGuards holding nothing

    JniUtfCharGuard() : utf_chars(nullptr) {}

    /// Release the held char sequence if there is one.

    ~JniUtfCharGuard() {

        if (utf_chars != nullptr) env->ReleaseStringUTFChars(jstr, utf_chars);

    }

    /// Try to get chars from jstr. If error is returned, utf_chars and get() remain

    /// to be nullptr, otherwise they point to a valid char sequence. The char sequence

    /// lives as long as this guard. jstr should not be null.

    static Status create(JNIEnv* env, jstring jstr, JniUtfCharGuard* out);

    /// Get the char sequence. Returns nullptr if the guard does hold a char sequence.

    const char* get() { return utf_chars; }

private:

    JNIEnv* env = nullptr;

    jstring jstr;

    const char* utf_chars = nullptr;

    DISALLOW_COPY_AND_ASSIGN(JniUtfCharGuard);

};

class JniLocalFrame {

public:

    JniLocalFrame() : env_(nullptr) {}

    ~JniLocalFrame() {

        if (env_ != nullptr) env_->PopLocalFrame(nullptr);

    }

    JniLocalFrame(JniLocalFrame&& other) noexcept : env_(other.env_) { other.env_ = nullptr; }

    /// Pushes a new JNI local frame. The frame can support max_local_ref local references.

    /// The number of local references created inside the frame might exceed max_local_ref,

    /// but there is no guarantee that memory will be available.

    /// Push should be called at most once.

    Status push(JNIEnv* env, int max_local_ref = 10) WARN_UNUSED_RESULT;

private:

    DISALLOW_COPY_AND_ASSIGN(JniLocalFrame);

    JNIEnv* env_ = nullptr;

};

template <class T>

Status SerializeThriftMsg(JNIEnv* env, T* msg, jbyteArray* serialized_msg) {

    int buffer_size = 100 * 1024; // start out with 100KB

    ThriftSerializer serializer(false, buffer_size);

    uint8_t* buffer = nullptr;

    uint32_t size = 0;

    RETURN_IF_ERROR(serializer.serialize(msg, &size, &buffer));

    // Make sure that 'size' is within the limit of INT_MAX as the use of

    // 'size' below takes int.

    if (size > INT_MAX) {

        return Status::InternalError(

                "The length of the serialization buffer ({} bytes) exceeds the limit of {} bytes",

                size, INT_MAX);

    }

    /// create jbyteArray given buffer

    *serialized_msg = env->NewByteArray(size);

    RETURN_ERROR_IF_EXC(env);

    if (*serialized_msg == NULL) return Status::InternalError("couldn't construct jbyteArray");

    env->SetByteArrayRegion(*serialized_msg, 0, size, reinterpret_cast<jbyte*>(buffer));

    RETURN_ERROR_IF_EXC(env);

    return Status::OK();

}

Unexecuted instantiation: _ZN5doris18SerializeThriftMsgINS_23TJdbcExecutorCtorParamsEEENS_6StatusEP7JNIEnv_PT_PP11_jbyteArray

Unexecuted instantiation: _ZN5doris18SerializeThriftMsgINS_26TJavaUdfExecutorCtorParamsEEENS_6StatusEP7JNIEnv_PT_PP11_jbyteArray

} // namespace doris