// 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 <fmt/format.h>

#include <gen_cpp/Types_types.h>

#include <jni.h>

#include <stdint.h>

#include <map>

#include <string>

#include <vector>

#include "common/status.h"

#include "core/data_type/data_type.h"

#include "exec/table_connector.h"

#include "exprs/aggregate/aggregate_function.h"

#include "util/jni-util.h"

namespace doris {

class RuntimeState;

class SlotDescriptor;

class TupleDescriptor;

class Block;

class IColumn;

class VExprContext;

struct JdbcConnectorParam {

    // use -1 as default value to find error earlier.

    int64_t catalog_id = -1;

    std::string driver_path;

    std::string driver_class;

    std::string resource_name;

    std::string driver_checksum;

    std::string jdbc_url;

    std::string user;

    std::string passwd;

    std::string query_string;

    std::string table_name;

    bool use_transaction = false;

    TOdbcTableType::type table_type;

    bool is_tvf = false;

    int32_t connection_pool_min_size = -1;

    int32_t connection_pool_max_size = -1;

    int32_t connection_pool_max_wait_time = -1;

    int32_t connection_pool_max_life_time = -1;

    bool connection_pool_keep_alive = false;

    const TupleDescriptor* tuple_desc = nullptr;

};

class JdbcConnector : public TableConnector {

public:

    struct JdbcStatistic {

        int64_t _load_jar_timer = 0;

        int64_t _init_connector_timer = 0;

        int64_t _get_data_timer = 0;

        int64_t _read_and_fill_vector_table_timer = 0;

        int64_t _jni_setup_timer = 0;

        int64_t _has_next_timer = 0;

        int64_t _prepare_params_timer = 0;

        int64_t _fill_block_timer = 0;

        int64_t _cast_timer = 0;

        int64_t _check_type_timer = 0;

        int64_t _execte_read_timer = 0;

        int64_t _connector_close_timer = 0;

    };

    JdbcConnector(const JdbcConnectorParam& param);

    ~JdbcConnector() override;

    Status open(RuntimeState* state, bool read = false);

    Status query() override;

    Status get_next(bool* eos, Block* block, int batch_size);

    Status append(Block* block, const VExprContextSPtrs& _output_vexpr_ctxs,

                  uint32_t start_send_row, uint32_t* num_rows_sent,

                  TOdbcTableType::type table_type = TOdbcTableType::MYSQL) override;

    Status exec_stmt_write(Block* block, const VExprContextSPtrs& output_vexpr_ctxs,

                           uint32_t* num_rows_sent) override;

    // use in JDBC transaction

    Status begin_trans() override; // should be call after connect and before query or init_to_write

    Status abort_trans() override; // should be call after transaction abort

    Status finish_trans() override; // should be call after transaction commit

    Status init_to_write(doris::RuntimeProfile* operator_profile) override {

        init_profile(operator_profile);

        return Status::OK();

    }

    JdbcStatistic& get_jdbc_statistic() { return _jdbc_statistic; }

    Status close(Status s = Status::OK()) override;

    Status test_connection();

    Status clean_datasource();

protected:

    JdbcConnectorParam _conn_param;

private:

    Status _register_func_id(JNIEnv* env);

    Status _get_reader_params(Block* block, JNIEnv* env, size_t column_size, Jni::LocalObject* ans);

    Status _cast_string_to_special(Block* block, JNIEnv* env, size_t column_size);

    Status _cast_string_to_hll(const SlotDescriptor* slot_desc, Block* block, int column_index,

                               int rows);

    Status _cast_string_to_bitmap(const SlotDescriptor* slot_desc, Block* block, int column_index,

                                  int rows);

    Status _cast_string_to_json(const SlotDescriptor* slot_desc, Block* block, int column_index,

                                int rows);

    Status _get_java_table_type(JNIEnv* env, TOdbcTableType::type table_type,

                                Jni::LocalObject* java_enum_obj);

    Status _get_real_url(const std::string& url, std::string* result_url);

    Status _check_and_return_default_driver_url(const std::string& url, std::string* result_url);

    bool _closed = false;

    Jni::GlobalClass _executor_factory_clazz;

    Jni::GlobalClass _executor_clazz;

    Jni::GlobalObject _executor_obj;

    Jni::MethodId _executor_factory_ctor_id;

    Jni::MethodId _executor_ctor_id;

    Jni::MethodId _executor_stmt_write_id;

    Jni::MethodId _executor_read_id;

    Jni::MethodId _executor_has_next_id;

    Jni::MethodId _executor_get_block_address_id;

    Jni::MethodId _executor_block_rows_id;

    Jni::MethodId _executor_close_id;

    Jni::MethodId _executor_begin_trans_id;

    Jni::MethodId _executor_finish_trans_id;

    Jni::MethodId _executor_abort_trans_id;

    Jni::MethodId _executor_test_connection_id;

    Jni::MethodId _executor_clean_datasource_id;

    std::map<int, int> _map_column_idx_to_cast_idx_hll;

    std::vector<DataTypePtr> _input_hll_string_types;

    std::map<int, int> _map_column_idx_to_cast_idx_bitmap;

    std::vector<DataTypePtr> _input_bitmap_string_types;

    std::map<int, int> _map_column_idx_to_cast_idx_json;

    std::vector<DataTypePtr> _input_json_string_types;

    JdbcStatistic _jdbc_statistic;

};

} // namespace doris