// This file is made available under Elastic License 2.0.
// This file is based on code available under the Apache license here:
//   https://github.com/apache/incubator-doris/blob/master/be/src/exec/es/es_predicate.cpp

// 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 "exec/es/es_predicate.h"

#include <gutil/strings/substitute.h>

#include <boost/algorithm/string.hpp>
#include <cstdint>
#include <map>
#include <sstream>

#include "common/logging.h"
#include "common/status.h"
#include "exec/es/es_query_builder.h"
#include "exprs/expr.h"
#include "exprs/expr_context.h"
#include "exprs/vectorized/column_ref.h"
#include "exprs/vectorized/in_const_predicate.hpp"
#include "gen_cpp/PlanNodes_types.h"
#include "runtime/client_cache.h"
#include "runtime/datetime_value.h"
#include "runtime/large_int_value.h"
#include "runtime/primitive_type.h"
#include "runtime/runtime_state.h"
#include "runtime/string_value.h"
#include "service/backend_options.h"
#include "storage/olap_common.h"
#include "storage/utils.h"
#include "util/debug_util.h"
#include "util/runtime_profile.h"

namespace starrocks {

using namespace std;

#define RETURN_ERROR_IF_EXPR_IS_NOT_SLOTREF(expr)                                          \
    do {                                                                                   \
        const Expr* expr_without_cast = Expr::expr_without_cast(expr);                     \
        if (expr_without_cast->node_type() != TExprNodeType::SLOT_REF) {                   \
            return Status::InternalError("build disjuncts failed: child is not slot ref"); \
        }                                                                                  \
    } while (false)

std::string SExtLiteral::value_to_string() {
    std::stringstream ss;
    switch (_type) {
    case TYPE_TINYINT:
        ss << std::to_string(get_byte());
        break;
    case TYPE_SMALLINT:
        ss << std::to_string(get_short());
        break;
    case TYPE_INT:
        ss << std::to_string(get_int());
        break;
    case TYPE_BIGINT:
        ss << std::to_string(get_long());
        break;
    case TYPE_FLOAT:
        ss << std::to_string(get_float());
        break;
    case TYPE_DOUBLE:
        ss << std::to_string(get_double());
        break;
    case TYPE_CHAR:
    case TYPE_VARCHAR:
        ss << get_string();
        break;
    case TYPE_DATE:
    case TYPE_DATETIME:
        ss << get_date_string();
        break;
    case TYPE_BOOLEAN:
        ss << std::to_string(get_bool());
        break;
    case TYPE_DECIMAL:
        ss << get_decimal_string();
        break;
    case TYPE_DECIMALV2:
        ss << get_decimalv2_string();
        break;
    case TYPE_LARGEINT:
        ss << get_largeint_string();
        break;
    default:
        DCHECK(false);
        break;
    }
    return ss.str();
}

template <typename T, typename... Ts>
static constexpr bool is_type_in() {
    return (std::is_same_v<T, Ts> || ...);
}

std::string VExtLiteral::_value_to_string(ColumnPtr& column) {
    auto v = column->get(0);
    std::string res;
    v.visit([&](auto& variant) {
        std::visit(
                [&](auto&& arg) {
                    using T = std::decay_t<decltype(arg)>;
                    if constexpr (is_type_in<T, std::monostate, int96_t, decimal12_t, DecimalV2Value,
                                             vectorized::DatumArray, HyperLogLog*, BitmapValue*, PercentileValue*>()) {
                        // ignore these types
                    } else if constexpr (std::is_same_v<T, Slice>) {
                        res = std::string(arg.data, arg.size);
                    } else if constexpr (std::is_same_v<T, __int128_t>) {
                        res = LargeIntValue::to_string(arg);
                    } else {
                        res = std::to_string(arg);
                    }
                },
                variant);
    });
    return res;
}

SExtLiteral::~SExtLiteral() = default;

int8_t SExtLiteral::get_byte() {
    DCHECK(_type == TYPE_TINYINT);
    return *(reinterpret_cast<int8_t*>(_value));
}

int16_t SExtLiteral::get_short() {
    DCHECK(_type == TYPE_SMALLINT);
    return *(reinterpret_cast<int16_t*>(_value));
}

int32_t SExtLiteral::get_int() {
    DCHECK(_type == TYPE_INT);
    return *(reinterpret_cast<int32_t*>(_value));
}

int64_t SExtLiteral::get_long() {
    DCHECK(_type == TYPE_BIGINT);
    return *(reinterpret_cast<int64_t*>(_value));
}

float SExtLiteral::get_float() {
    DCHECK(_type == TYPE_FLOAT);
    return *(reinterpret_cast<float*>(_value));
}

double SExtLiteral::get_double() {
    DCHECK(_type == TYPE_DOUBLE);
    return *(reinterpret_cast<double*>(_value));
}

std::string SExtLiteral::get_string() {
    DCHECK(_type == TYPE_VARCHAR || _type == TYPE_CHAR);
    return (reinterpret_cast<StringValue*>(_value))->to_string();
}

std::string SExtLiteral::get_date_string() {
    DCHECK(_type == TYPE_DATE || _type == TYPE_DATETIME);
    DateTimeValue date_value = *reinterpret_cast<DateTimeValue*>(_value);
    if (_type == TYPE_DATE) {
        date_value.cast_to_date();
    }

    char str[MAX_DTVALUE_STR_LEN];
    date_value.to_string(str);
    return std::string(str, strlen(str));
}

bool SExtLiteral::get_bool() {
    DCHECK(_type == TYPE_BOOLEAN);
    return *(reinterpret_cast<bool*>(_value));
}

std::string SExtLiteral::get_decimal_string() {
    DCHECK(_type == TYPE_DECIMAL);
    return reinterpret_cast<DecimalValue*>(_value)->to_string();
}

std::string SExtLiteral::get_decimalv2_string() {
    DCHECK(_type == TYPE_DECIMALV2);
    return reinterpret_cast<DecimalV2Value*>(_value)->to_string();
}

std::string SExtLiteral::get_largeint_string() {
    DCHECK(_type == TYPE_LARGEINT);
    __int128 v = unaligned_load<__int128>(_value);
    return LargeIntValue::to_string(v);
}

EsPredicate::EsPredicate(ExprContext* context, const TupleDescriptor* tuple_desc, ObjectPool* pool)
        : _context(context), _tuple_desc(tuple_desc), _es_query_status(Status::OK()), _pool(pool) {}

EsPredicate::~EsPredicate() {
    for (auto& _disjunct : _disjuncts) {
        delete _disjunct;
    }
    _disjuncts.clear();
}

Status EsPredicate::build_disjuncts_list(bool use_vectorized) {
    return _vec_build_disjuncts_list(_context->root());
}

// make sure to build by build_disjuncts_list
const std::vector<ExtPredicate*>& EsPredicate::get_predicate_list() {
    return _disjuncts;
}

static bool ignore_cast(const SlotDescriptor* slot, const Expr* expr) {
    if (slot->type().is_date_type() && expr->type().is_date_type()) {
        return true;
    }
    if (slot->type().is_string_type() && expr->type().is_string_type()) {
        return true;
    }
    return false;
}

static bool is_literal_node(const Expr* expr) {
    switch (expr->node_type()) {
    case TExprNodeType::BOOL_LITERAL:
    case TExprNodeType::INT_LITERAL:
    case TExprNodeType::LARGE_INT_LITERAL:
    case TExprNodeType::FLOAT_LITERAL:
    case TExprNodeType::DECIMAL_LITERAL:
    case TExprNodeType::STRING_LITERAL:
    case TExprNodeType::DATE_LITERAL:
        return true;
    default:
        return false;
    }
}

Status EsPredicate::_vec_build_disjuncts_list(const Expr* conjunct) {
    bool handled = false;
    RETURN_IF_ERROR(_build_binary_predicate(conjunct, &handled));
    RETURN_IF_ERROR(_build_functioncall_predicate(conjunct, &handled));
    RETURN_IF_ERROR(_build_in_predicate(conjunct, &handled));
    RETURN_IF_ERROR(_build_compound_predicate(conjunct, &handled));
    if (handled) {
        return Status::OK();
    } else {
        return Status::InternalError(
                fmt::format("build disjuncts failed: node type {} is not supported", conjunct->node_type()));
    }
}

Status EsPredicate::_build_binary_predicate(const Expr* conjunct, bool* handled) {
    using ColumnRef = vectorized::ColumnRef;
    if (TExprNodeType::BINARY_PRED == conjunct->node_type()) {
        DCHECK_EQ(conjunct->children().size(), 2);
        *handled = true;

        ColumnRef* column_ref = nullptr;
        Expr* expr = nullptr;
        TExprOpcode::type op;
        // k1 = 2  k1 is float (marked for processing later),
        // doris on es should ignore this doris native cast transformation, we push down this `cast` to elasticsearch
        // conjunct->get_child(0)->node_type() return CAST_EXPR
        // conjunct->get_child(1)->node_type() return FLOAT_LITERAL
        // the left child is literal and right child is SlotRef maybe not happend, but here we just process
        // this situation regardless of the rewrite logic from the FE's Query Engine
        if (TExprNodeType::SLOT_REF == conjunct->get_child(0)->node_type() ||
            TExprNodeType::CAST_EXPR == conjunct->get_child(0)->node_type()) {
            expr = conjunct->get_child(1);
            // process such as sub-query: select * from (select split_part(k, "_", 1) as new_field from table) t where t.new_field > 1;
            RETURN_ERROR_IF_EXPR_IS_NOT_SLOTREF(conjunct->get_child(0));
            // process cast expr, such as:
            // k (float) > 2.0, k(int) > 3.2
            column_ref = const_cast<ColumnRef*>(
                    down_cast<const ColumnRef*>(Expr::expr_without_cast(conjunct->get_child(0))));
            op = conjunct->op();
        } else if (TExprNodeType::SLOT_REF == conjunct->get_child(1)->node_type() ||
                   TExprNodeType::CAST_EXPR == conjunct->get_child(1)->node_type()) {
            expr = conjunct->get_child(0);
            RETURN_ERROR_IF_EXPR_IS_NOT_SLOTREF(conjunct->get_child(1));
            column_ref = const_cast<ColumnRef*>(
                    down_cast<const ColumnRef*>(Expr::expr_without_cast(conjunct->get_child(1))));
            op = conjunct->op();
        } else {
            return Status::InternalError("build disjuncts failed: no SLOT_REF child");
        }

        const SlotDescriptor* slot_desc = get_slot_desc(column_ref->slot_id());
        if (column_ref == nullptr) {
            return Status::InternalError("build disjuncts failed: slot_desc is null");
        }

        if (!is_literal_node(expr)) {
            return Status::InternalError("build disjuncts failed: expr is not literal type");
        }

        // how to process literal
        auto literal = _pool->add(new VExtLiteral(expr->type().type, _context->evaluate(expr, nullptr)));
        std::string col = slot_desc->col_name();

        // ES does not support non-bool literal pushdown for bool type
        if (column_ref->type().type == TYPE_BOOLEAN && expr->type().type != TYPE_BOOLEAN) {
            return Status::InternalError("ES does not support non-bool literal pushdown");
        }

        if (_field_context.find(col) != _field_context.end()) {
            col = _field_context[col];
        }
        ExtPredicate* predicate =
                new ExtBinaryPredicate(TExprNodeType::BINARY_PRED, col, slot_desc->type(), op, literal);

        _disjuncts.push_back(predicate);
    }
    return Status::OK();
}

Status EsPredicate::_build_functioncall_predicate(const Expr* conjunct, bool* handled) {
    using ColumnRef = vectorized::ColumnRef;

    if (TExprNodeType::FUNCTION_CALL == conjunct->node_type()) {
        *handled = true;

        std::string fname = conjunct->fn().name.function_name;
        if (fname == "esquery") {
            if (conjunct->children().size() != 2) {
                return Status::InternalError("build disjuncts failed: number of childs is not 2");
            }
            Expr* expr = conjunct->get_child(1);
            auto literal = _pool->add(new VExtLiteral(expr->type().type, _context->evaluate(expr, nullptr)));
            std::vector<ExtLiteral*> query_conditions;
            query_conditions.emplace_back(literal);
            std::vector<ExtColumnDesc> cols;
            ExtPredicate* predicate = new ExtFunction(TExprNodeType::FUNCTION_CALL, "esquery", cols, query_conditions);
            if (_es_query_status.ok()) {
                _es_query_status = BooleanQueryBuilder::check_es_query(*(ExtFunction*)predicate);
                if (!_es_query_status.ok()) {
                    delete predicate;
                    return _es_query_status;
                }
            }
            _disjuncts.push_back(predicate);
        } else if (fname == "is_null_pred" || fname == "is_not_null_pred") {
            if (conjunct->children().size() != 1) {
                return Status::InternalError("build disjuncts failed: number of childs is not 1");
            }
            // such as sub-query: select * from (select split_part(k, "_", 1) as new_field from table) t where t.new_field > 1;
            // conjunct->get_child(0)->node_type() == TExprNodeType::FUNCTION_CALL, at present doris on es can not support push down function
            RETURN_ERROR_IF_EXPR_IS_NOT_SLOTREF(conjunct->get_child(0));

            ColumnRef* column_ref = const_cast<ColumnRef*>(
                    down_cast<const ColumnRef*>(Expr::expr_without_cast(conjunct->get_child(0))));

            const SlotDescriptor* slot_desc = get_slot_desc(column_ref->slot_id());

            if (slot_desc == nullptr) {
                return Status::InternalError("build disjuncts failed: no SLOT_REF child");
            }
            bool is_not_null = fname == "is_not_null_pred" ? true : false;
            std::string col = slot_desc->col_name();
            if (_field_context.find(col) != _field_context.end()) {
                col = _field_context[col];
            }
            // use TExprNodeType::IS_NULL_PRED for BooleanQueryBuilder translate
            ExtIsNullPredicate* predicate =
                    new ExtIsNullPredicate(TExprNodeType::IS_NULL_PRED, col, slot_desc->type(), is_not_null);
            _disjuncts.push_back(predicate);
        } else if (fname == "like") {
            if (conjunct->children().size() != 2) {
                return Status::InternalError("build disjuncts failed: number of childs is not 2");
            }
            ColumnRef* column_ref = nullptr;
            Expr* expr = nullptr;
            if (TExprNodeType::SLOT_REF == conjunct->get_child(0)->node_type()) {
                expr = conjunct->get_child(1);
                column_ref = const_cast<ColumnRef*>(down_cast<const ColumnRef*>(conjunct->get_child(0)));
            } else if (TExprNodeType::SLOT_REF == conjunct->get_child(1)->node_type()) {
                expr = conjunct->get_child(0);
                column_ref = const_cast<ColumnRef*>(down_cast<const ColumnRef*>(conjunct->get_child(1)));
            } else {
                return Status::InternalError("build disjuncts failed: no SLOT_REF child");
            }
            const SlotDescriptor* slot_desc = get_slot_desc(column_ref->slot_id());
            if (slot_desc == nullptr) {
                return Status::InternalError("build disjuncts failed: slot_desc is null");
            }

            PrimitiveType type = expr->type().type;
            if (type != TYPE_VARCHAR && type != TYPE_CHAR) {
                return Status::InternalError("build disjuncts failed: like value is not a string");
            }
            std::string col = slot_desc->col_name();
            if (_field_context.find(col) != _field_context.end()) {
                col = _field_context[col];
            }

            auto literal = _pool->add(new VExtLiteral(type, _context->evaluate(expr, nullptr)));
            ExtPredicate* predicate = new ExtLikePredicate(TExprNodeType::LIKE_PRED, col, slot_desc->type(), literal);

            _disjuncts.push_back(predicate);
        } else {
            std::stringstream ss;
            ss << "can not process function predicate[ " << fname << " ]";
            return Status::InternalError(ss.str());
        }
        return Status::OK();
    }

    return Status::OK();
}

template <PrimitiveType type, typename Func>
Status build_inpred_values(const Predicate* pred, bool& is_not_in, Func&& func) {
    const auto* vpred = down_cast<const vectorized::VectorizedInConstPredicate<type>*>(pred);
    const auto& hash_set = vpred->hash_set();
    bool has_null = vpred->null_in_set();
    is_not_in = vpred->is_not_in();
    if (has_null) {
        return Status::InternalError("build disjuncts failed: hash set has a null value");
    }
    for (auto v : hash_set) {
        func(v);
    }
    return Status::OK();
}

#define BUILD_INPRED_VALUES(TYPE)                                                                         \
    case TYPE: {                                                                                          \
        RETURN_IF_ERROR(build_inpred_values<TYPE>(pred, is_not_in, [&](auto& v) {                         \
            in_pred_values.emplace_back(new SExtLiteral(slot_desc->type().type, static_cast<void*>(&v))); \
        }));                                                                                              \
        break;                                                                                            \
    }

Status EsPredicate::_build_in_predicate(const Expr* conjunct, bool* handled) {
    using ColumnRef = vectorized::ColumnRef;

    if (TExprNodeType::IN_PRED == conjunct->node_type()) {
        *handled = true;
        // the op code maybe FILTER_NEW_IN, it means there is function in list
        // like col_a in (abs(1))
        if (TExprOpcode::FILTER_IN != conjunct->op() && TExprOpcode::FILTER_NOT_IN != conjunct->op()) {
            return Status::InternalError(
                    "build disjuncts failed: "
                    "opcode in IN_PRED is neither FILTER_IN nor FILTER_NOT_IN");
        }

        std::vector<ExtLiteral*> in_pred_values;
        const Predicate* pred = static_cast<const Predicate*>(conjunct);

        const Expr* expr = Expr::expr_without_cast(pred->get_child(0));
        if (expr->node_type() != TExprNodeType::SLOT_REF) {
            return Status::InternalError("build disjuncts failed: node type is not slot ref");
        }

        const SlotDescriptor* slot_desc = get_slot_desc(static_cast<const ColumnRef*>(expr)->slot_id());
        if (slot_desc == nullptr) {
            return Status::InternalError("build disjuncts failed: slot_desc is null");
        }

        if (pred->get_child(0)->type().type != slot_desc->type().type) {
            if (!ignore_cast(slot_desc, pred->get_child(0))) {
                return Status::InternalError("build disjuncts failed");
            }
        }

        bool is_not_in = false;
        // insert in list to ExtLiteral
        switch (expr->type().type) {
            BUILD_INPRED_VALUES(TYPE_BOOLEAN);
            BUILD_INPRED_VALUES(TYPE_INT);
            BUILD_INPRED_VALUES(TYPE_TINYINT);
            BUILD_INPRED_VALUES(TYPE_SMALLINT);
            BUILD_INPRED_VALUES(TYPE_BIGINT);
            BUILD_INPRED_VALUES(TYPE_LARGEINT);
            BUILD_INPRED_VALUES(TYPE_FLOAT);
            BUILD_INPRED_VALUES(TYPE_DOUBLE);
        case TYPE_DATE: {
            RETURN_IF_ERROR(build_inpred_values<TYPE_DATE>(pred, is_not_in, [&](auto& v) {
                // declare_type(v) == DateValue
                int year, month, day;
                v.to_date(&year, &month, &day);
                DateTimeValue date_value(TIME_DATE, year, month, day, 0, 0, 0, 0);
                in_pred_values.emplace_back(new SExtLiteral(slot_desc->type().type, static_cast<void*>(&date_value)));
            }));
            break;
        }
        case TYPE_DATETIME: {
            RETURN_IF_ERROR(build_inpred_values<TYPE_DATETIME>(pred, is_not_in, [&](auto& v) {
                // declare_type(v) == TimestampValue
                int year, month, day, hour, minute, second, usec;
                v.to_timestamp(&year, &month, &day, &hour, &minute, &second, &usec);
                DateTimeValue ts_value(TIME_DATETIME, year, month, day, hour, minute, second, usec);
                in_pred_values.emplace_back(new SExtLiteral(slot_desc->type().type, static_cast<void*>(&ts_value)));
            }));
            break;
        }
        case TYPE_CHAR:
        case TYPE_VARCHAR: {
            RETURN_IF_ERROR(build_inpred_values<TYPE_VARCHAR>(pred, is_not_in, [&](auto& v) {
                // declare_type(v) == Slice
                StringValue string_val(v.data, v.size);
                in_pred_values.emplace_back(new SExtLiteral(slot_desc->type().type, static_cast<void*>(&string_val)));
            }));
            break;
        }
        default:
            DCHECK(false) << "unsupport type:" << expr->type().type;
            return Status::InternalError("unsupport type to push down to ES");
        }

        std::string col = slot_desc->col_name();
        if (_field_context.find(col) != _field_context.end()) {
            col = _field_context[col];
        }
        ExtPredicate* predicate = new ExtInPredicate(TExprNodeType::IN_PRED, is_not_in, col, slot_desc->type(),
                                                     std::move(in_pred_values));
        _disjuncts.push_back(predicate);

        return Status::OK();
    }
    return Status::OK();
}

Status EsPredicate::_build_compound_predicate(const Expr* conjunct, bool* handled) {
    using ColumnRef = vectorized::ColumnRef;

    if (TExprNodeType::COMPOUND_PRED == conjunct->node_type()) {
        *handled = true;
        // processe COMPOUND_AND, such as:
        // k = 1 or (k1 = 7 and (k2 in (6,7) or k3 = 12))
        // k1 = 7 and (k2 in (6,7) or k3 = 12) is compound pred, we should rebuild this sub tree
        if (conjunct->op() == TExprOpcode::COMPOUND_AND) {
            std::vector<EsPredicate*> conjuncts;
            for (int i = 0; i < conjunct->get_num_children(); ++i) {
                EsPredicate* predicate = _pool->add(new EsPredicate(_context, _tuple_desc, _pool));
                predicate->set_field_context(_field_context);
                Status status = predicate->_vec_build_disjuncts_list(conjunct->children()[i]);
                if (status.ok()) {
                    conjuncts.push_back(predicate);
                } else {
                    return Status::InternalError("build COMPOUND_AND conjuncts failed");
                }
            }
            ExtCompPredicates* compound_predicate = new ExtCompPredicates(TExprOpcode::COMPOUND_AND, conjuncts);
            _disjuncts.push_back(compound_predicate);
            return Status::OK();
        } else if (conjunct->op() == TExprOpcode::COMPOUND_NOT) {
            // reserved for processing COMPOUND_NOT
            return Status::InternalError("currently do not support COMPOUND_NOT push-down");
        }
        DCHECK(conjunct->op() == TExprOpcode::COMPOUND_OR);
        RETURN_IF_ERROR(_vec_build_disjuncts_list(conjunct->get_child(0)));
        RETURN_IF_ERROR(_vec_build_disjuncts_list(conjunct->get_child(1)));
    }
    return Status::OK();
}

const SlotDescriptor* EsPredicate::get_slot_desc(const SlotRef* slotRef) {
    return get_slot_desc(slotRef->slot_id());
}

const SlotDescriptor* EsPredicate::get_slot_desc(SlotId slot_id) {
    const SlotDescriptor* slot_desc = nullptr;
    for (SlotDescriptor* slot : _tuple_desc->slots()) {
        if (slot->id() == slot_id) {
            slot_desc = slot;
            break;
        }
    }
    return slot_desc;
}

} // namespace starrocks
