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

#include "exprs/function/function_rpc.h"

#include <brpc/controller.h>

#include <fmt/format.h>

#include <gen_cpp/function_service.pb.h>

#include <gen_cpp/types.pb.h>

#include <algorithm>

#include <memory>

#include <utility>

#include "common/status.h"

#include "core/column/column.h"

#include "core/data_type_serde/data_type_serde.h"

#include "runtime/exec_env.h"

#include "util/brpc_client_cache.h"

namespace doris {

#include "common/compile_check_begin.h"

RPCFnImpl::RPCFnImpl(const TFunction& fn) : _fn(fn) {

    _function_name = _fn.scalar_fn.symbol;

    _server_addr = _fn.hdfs_location;

    _client = ExecEnv::GetInstance()->brpc_function_client_cache()->get_client(_server_addr);

    _signature = fmt::format("{}: [{}/{}]", _fn.name.function_name, _fn.hdfs_location,

                             _fn.scalar_fn.symbol);

}

Status RPCFnImpl::vec_call(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

                           uint32_t result, size_t input_rows_count) {

    PFunctionCallRequest request;

    PFunctionCallResponse response;

    if (_client == nullptr) {

        return Status::InternalError(

                "call to rpc function {} failed: init rpc error, server addr = {}", _signature,

                _server_addr);

    }

    request.set_function_name(_function_name);

    RETURN_IF_ERROR(_convert_block_to_proto(block, arguments, input_rows_count, &request));

    brpc::Controller cntl;

    _client->fn_call(&cntl, &request, &response, nullptr);

    if (cntl.Failed()) {

        return Status::InternalError("call to rpc function {} failed: {}", _signature,

                                     cntl.ErrorText());

    }

    if (!response.has_status() || response.result_size() == 0) {

        return Status::InternalError("call rpc function {} failed: status or result is not set.",

                                     _signature);

    }

    if (response.status().status_code() != 0) {

        return Status::InternalError("call to rpc function {} failed: {}", _signature,

                                     response.status().DebugString());

    }

    RETURN_IF_ERROR(_convert_to_block(block, response.result(0), result));

    return Status::OK();

}

Status RPCFnImpl::_convert_block_to_proto(Block& block, const ColumnNumbers& arguments,

                                          size_t input_rows_count, PFunctionCallRequest* request) {

    size_t row_count = std::min(block.rows(), input_rows_count);

    for (size_t col_idx : arguments) {

        PValues* arg = request->add_args();

        ColumnWithTypeAndName& column = block.get_by_position(col_idx);

        arg->set_has_null(column.column->has_null(0, row_count));

        auto col = column.column->convert_to_full_column_if_const();

        RETURN_IF_ERROR(column.type->get_serde()->write_column_to_pb(*col, *arg, 0, row_count));

    }

    return Status::OK();

}

Status RPCFnImpl::_convert_to_block(Block& block, const PValues& result, size_t pos) {

    auto data_type = block.get_data_type(pos);

    auto col = data_type->create_column();

    auto serde = data_type->get_serde();

    RETURN_IF_ERROR(serde->read_column_from_pb(*col, result));

    block.replace_by_position(pos, std::move(col));

    return Status::OK();

}

FunctionRPC::FunctionRPC(const TFunction& fn, const DataTypes& argument_types,

                         const DataTypePtr& return_type)

        : _argument_types(argument_types), _return_type(return_type), _tfn(fn) {}

Status FunctionRPC::open(FunctionContext* context, FunctionContext::FunctionStateScope scope) {

    if (scope == FunctionContext::FRAGMENT_LOCAL) {

        std::shared_ptr<RPCFnImpl> fn = std::make_shared<RPCFnImpl>(_tfn);

        if (!fn->available()) {

            return Status::InternalError("rpc env init error");

        }

        context->set_function_state(FunctionContext::FRAGMENT_LOCAL, fn);

    }

    return Status::OK();

}

} // namespace doris