// 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 <glog/logging.h>

#include <algorithm>

#include <cstddef>

#include <memory>

#include <optional>

#include <unordered_map>

#include <utility>

#include <vector>

#include "common/factory_creator.h"

#include "common/status.h"

#include "core/block/block.h"

#include "core/block/column_with_type_and_name.h"

#include "core/column/column.h"

#include "exec/runtime_filter/runtime_filter_selectivity.h"

#include "exprs/function_context.h"

#include "exprs/vexpr_fwd.h"

#include "runtime/runtime_state.h"

#include "runtime/scan_filter_profile.h"

#include "storage/index/ann/ann_range_search_runtime.h"

#include "storage/index/ann/ann_search_params.h"

#include "storage/index/inverted/inverted_index_reader.h"

#include "storage/segment/column_reader.h"

namespace doris {

class RowDescriptor;

class RuntimeState;

} // namespace doris

namespace doris::segment_v2 {

class Segment;

class ColumnIterator;

} // namespace doris::segment_v2

namespace doris {

class ScoreRuntime;

using ScoreRuntimeSPtr = std::shared_ptr<ScoreRuntime>;

class IndexExecContext {

public:

    IndexExecContext(const std::vector<ColumnId>& col_ids,

                     const std::vector<std::unique_ptr<segment_v2::IndexIterator>>& index_iterators,

                     const std::vector<IndexFieldNameAndTypePair>& storage_name_and_type_vec,

                     std::unordered_map<ColumnId, std::unordered_map<const VExpr*, bool>>&

                             common_expr_index_status,

                     ScoreRuntimeSPtr score_runtime, segment_v2::Segment* segment,

                     const segment_v2::ColumnIteratorOptions& column_iter_opts)

            : _col_ids(col_ids),

              _index_iterators(index_iterators),

              _storage_name_and_type(storage_name_and_type_vec),

              _expr_index_status(common_expr_index_status),

              _score_runtime(std::move(score_runtime)),

              _segment(segment),

              _column_iter_opts(column_iter_opts) {}

    segment_v2::IndexIterator* get_inverted_index_iterator_by_column_id(int column_index) const {

        if (column_index < 0 || column_index >= _col_ids.size()) {

            return nullptr;

        }

        const auto& column_id = _col_ids[column_index];

        if (column_id >= _index_iterators.size()) {

            return nullptr;

        }

        if (!_index_iterators[column_id]) {

            return nullptr;

        }

        return _index_iterators[column_id].get();

    }

    segment_v2::IndexIterator* get_inverted_index_iterator_by_id(ColumnId column_id) const {

        if (column_id >= _index_iterators.size()) {

            return nullptr;

        }

        if (!_index_iterators[column_id]) {

            return nullptr;

        }

        return _index_iterators[column_id].get();

    }

    const IndexFieldNameAndTypePair* get_storage_name_and_type_by_column_id(

            int column_index) const {

        if (column_index < 0 || column_index >= _col_ids.size()) {

            return nullptr;

        }

        const auto& column_id = _col_ids[column_index];

        if (column_id >= _storage_name_and_type.size()) {

            return nullptr;

        }

        return &_storage_name_and_type[column_id];

    }

    const IndexFieldNameAndTypePair* get_storage_name_and_type_by_id(ColumnId column_id) const {

        if (column_id >= _storage_name_and_type.size()) {

            return nullptr;

        }

        return &_storage_name_and_type[column_id];

    }

    int column_index_by_id(ColumnId column_id) const {

        for (int i = 0; i < _col_ids.size(); ++i) {

            if (_col_ids[i] == column_id) {

                return i;

            }

        }

        return -1;

    }

    bool get_column_id(int column_index, ColumnId* column_id) const {

        if (column_id == nullptr) {

            return false;

        }

        if (column_index < 0 || column_index >= _col_ids.size()) {

            return false;

        }

        *column_id = _col_ids[column_index];

        return true;

    }

    segment_v2::Segment* segment() const { return _segment; }

    const segment_v2::ColumnIteratorOptions& column_iter_opts() const { return _column_iter_opts; }

    bool has_index_result_for_expr(const VExpr* expr) const {

        return _index_result_bitmap.contains(expr);

    }

    void set_index_result_for_expr(const VExpr* expr,

                                   segment_v2::InvertedIndexResultBitmap bitmap) {

        _index_result_bitmap[expr] = std::move(bitmap);

    }

    std::unordered_map<const VExpr*, segment_v2::InvertedIndexResultBitmap>&

    get_index_result_bitmap() {

        return _index_result_bitmap;

    }

    std::unordered_map<const VExpr*, ColumnPtr>& get_index_result_column() {

        return _index_result_column;

    }

    const segment_v2::InvertedIndexResultBitmap* get_index_result_for_expr(const VExpr* expr) {

        auto iter = _index_result_bitmap.find(expr);

        if (iter == _index_result_bitmap.end()) {

            return nullptr;

        }

        return &iter->second;

    }

    void set_index_result_column_for_expr(const VExpr* expr, ColumnPtr column) {

        _index_result_column[expr] = std::move(column);

    }

    void set_true_for_index_status(const VExpr* expr, int column_index) {

        if (column_index < 0 || column_index >= _col_ids.size()) {

            return;

        }

        const auto& column_id = _col_ids[column_index];

        if (_expr_index_status.contains(column_id)) {

            if (_expr_index_status[column_id].contains(expr)) {

                _expr_index_status[column_id][expr] = true;

            }

        }

    }

    ScoreRuntimeSPtr get_score_runtime() const { return _score_runtime; }

    void set_analyzer_ctx_for_expr(const VExpr* expr, InvertedIndexAnalyzerCtxSPtr analyzer_ctx) {

        if (expr == nullptr || analyzer_ctx == nullptr) {

            return;

        }

        _expr_analyzer_ctx[expr] = std::move(analyzer_ctx);

    }

    const InvertedIndexAnalyzerCtx* get_analyzer_ctx_for_expr(const VExpr* expr) const {

        auto iter = _expr_analyzer_ctx.find(expr);

        if (iter == _expr_analyzer_ctx.end()) {

            return nullptr;

        }

        return iter->second.get();

    }

    void set_index_query_context(segment_v2::IndexQueryContextPtr index_query_context) {

        _index_query_context = index_query_context;

    }

    const segment_v2::IndexQueryContextPtr& get_index_query_context() const {

        return _index_query_context;

    }

private:

    // A reference to a vector of column IDs for the current expression's output columns.

    const std::vector<ColumnId>& _col_ids;

    // A reference to a vector of unique pointers to index iterators.

    const std::vector<std::unique_ptr<segment_v2::IndexIterator>>& _index_iterators;

    // A reference to a vector of storage name and type pairs related to schema.

    const std::vector<IndexFieldNameAndTypePair>& _storage_name_and_type;

    // A map of expressions to their corresponding inverted index result bitmaps.

    std::unordered_map<const VExpr*, segment_v2::InvertedIndexResultBitmap> _index_result_bitmap;

    // A map of expressions to their corresponding result columns.

    std::unordered_map<const VExpr*, ColumnPtr> _index_result_column;

    // Per-expression analyzer context for inverted index evaluation.

    std::unordered_map<const VExpr*, InvertedIndexAnalyzerCtxSPtr> _expr_analyzer_ctx;

    // A reference to a map of common expressions to their inverted index evaluation status.

    std::unordered_map<ColumnId, std::unordered_map<const VExpr*, bool>>& _expr_index_status;

    ScoreRuntimeSPtr _score_runtime;

    segment_v2::Segment* _segment = nullptr; // Ref

    segment_v2::ColumnIteratorOptions _column_iter_opts;

    segment_v2::IndexQueryContextPtr _index_query_context;

};

class VExprContext {

    ENABLE_FACTORY_CREATOR(VExprContext);

public:

    VExprContext(VExprSPtr expr) : _root(std::move(expr)) {}

    ~VExprContext();

    [[nodiscard]] Status prepare(RuntimeState* state, const RowDescriptor& row_desc);

    [[nodiscard]] Status open(RuntimeState* state);

    [[nodiscard]] Status clone(RuntimeState* state, VExprContextSPtr& new_ctx);

    [[nodiscard]] Status execute(Block* block, int* result_column_id);

    [[nodiscard]] Status execute(const Block* block, ColumnPtr& result_column);

    [[nodiscard]] Status execute(const Block* block, ColumnWithTypeAndName& result_data);

    [[nodiscard]] DataTypePtr execute_type(const Block* block);

    [[nodiscard]] const std::string& expr_name() const;

    [[nodiscard]] bool is_blockable() const;

    [[nodiscard]] Status execute_const_expr(ColumnWithTypeAndName& result);

    double execute_cost() const;

    VExprSPtr root() { return _root; }

    void set_root(const VExprSPtr& expr) { _root = expr; }

    void set_index_context(std::shared_ptr<IndexExecContext> index_context) {

        _index_context = std::move(index_context);

    }

    std::shared_ptr<IndexExecContext> get_index_context() const { return _index_context; }

    void attach_scan_filter(ScanFilterHandle handle) { _scan_filter_handle = std::move(handle); }

    const ScanFilterHandle& scan_filter_handle() const { return _scan_filter_handle; }

    /// Creates a FunctionContext, and returns the index that's passed to fn_context() to

    /// retrieve the created context. Exprs that need a FunctionContext should call this in

    /// Prepare() and save the returned index. 'varargs_buffer_size', if specified, is the

    /// size of the varargs buffer in the created FunctionContext (see udf-internal.h).

    int register_function_context(RuntimeState* state, const DataTypePtr& return_type,

                                  const std::vector<DataTypePtr>& arg_types);

    /// Retrieves a registered FunctionContext. 'i' is the index returned by the call to

    /// register_function_context(). This should only be called by VExprs.

    FunctionContext* fn_context(int i) {

        if (i < 0 || i >= _fn_contexts.size()) {

            throw Exception(ErrorCode::INTERNAL_ERROR,

                            "fn_context index invalid, index={}, _fn_contexts.size()={}", i,

                            _fn_contexts.size());

        }

        return _fn_contexts[i].get();

    }

    // execute expr with inverted index which column a, b has inverted indexes

    //  but some situation although column b has indexes, but apply index is not useful, we should

    //  skip this expr, just do not apply index anymore.

    [[nodiscard]] Status evaluate_inverted_index(uint32_t segment_num_rows);

    bool all_expr_inverted_index_evaluated();

    Status execute_filter(const Block* block, uint8_t* __restrict result_filter_data, size_t rows,

                          bool accept_null, bool* can_filter_all);

    [[nodiscard]] static Status filter_block(VExprContext* vexpr_ctx, Block* block);

    [[nodiscard]] static Status filter_block(

            const VExprContextSPtrs& expr_contexts, Block* block, size_t column_to_keep,

            std::optional<ScanFilterStage> scan_filter_stage = std::nullopt);

    [[nodiscard]] static Status execute_conjuncts(

            const VExprContextSPtrs& ctxs, const std::vector<IColumn::Filter*>* filters,

            bool accept_null, const Block* block, IColumn::Filter* result_filter,

            bool* can_filter_all, std::optional<ScanFilterStage> scan_filter_stage = std::nullopt);

    [[nodiscard]] static Status execute_conjuncts(const VExprContextSPtrs& conjuncts,

                                                  const Block* block, ColumnUInt8& null_map,

                                                  IColumn::Filter& result_filter);

    static Status execute_conjuncts(

            const VExprContextSPtrs& ctxs, const std::vector<IColumn::Filter*>* filters,

            Block* block, IColumn::Filter* result_filter, bool* can_filter_all,

            std::optional<ScanFilterStage> scan_filter_stage = std::nullopt);

    [[nodiscard]] static Status execute_conjuncts_and_filter_block(

            const VExprContextSPtrs& ctxs, Block* block, std::vector<uint32_t>& columns_to_filter,

            int column_to_keep, std::optional<ScanFilterStage> scan_filter_stage = std::nullopt);

    static Status execute_conjuncts_and_filter_block(

            const VExprContextSPtrs& ctxs, Block* block, std::vector<uint32_t>& columns_to_filter,

            int column_to_keep, IColumn::Filter& filter,

            std::optional<ScanFilterStage> scan_filter_stage = std::nullopt);

    [[nodiscard]] static Status get_output_block_after_execute_exprs(const VExprContextSPtrs&,

                                                                     const Block&, Block*,

                                                                     bool do_projection = false);

    int get_last_result_column_id() const {

        DCHECK(_last_result_column_id != -1);

        return _last_result_column_id;

    }

    RuntimeFilterSelectivity& get_runtime_filter_selectivity() {

        if (!_rf_selectivity) {

            throw Exception(ErrorCode::INTERNAL_ERROR, "RuntimeFilterSelectivity is null");

        }

        return *_rf_selectivity;

    }

    FunctionContext::FunctionStateScope get_function_state_scope() const {

        return _is_clone ? FunctionContext::THREAD_LOCAL : FunctionContext::FRAGMENT_LOCAL;

    }

    void clone_fn_contexts(VExprContext* other);

    VExprContext& operator=(const VExprContext& other) {

        if (this == &other) {

            return *this;

        }

        _root = other._root;

        _is_clone = other._is_clone;

        _prepared = other._prepared;

        _opened = other._opened;

        for (const auto& fn : other._fn_contexts) {

            _fn_contexts.emplace_back(fn->clone());

        }

        _last_result_column_id = other._last_result_column_id;

        _depth_num = other._depth_num;

        _scan_filter_handle = other._scan_filter_handle;

        return *this;

    }

    VExprContext& operator=(VExprContext&& other) {

        _root = other._root;

        other._root = nullptr;

        _is_clone = other._is_clone;

        _prepared = other._prepared;

        _opened = other._opened;

        _fn_contexts = std::move(other._fn_contexts);

        _last_result_column_id = other._last_result_column_id;

        _depth_num = other._depth_num;

        _scan_filter_handle = other._scan_filter_handle;

        return *this;

    }

    [[nodiscard]] static size_t get_memory_usage(const VExprContextSPtrs& contexts) {

        size_t usage = 0;

        std::for_each(contexts.cbegin(), contexts.cend(),

                      [&usage](auto&& context) { usage += context->_memory_usage; });

        return usage;

    }

    [[nodiscard]] size_t get_memory_usage() const { return _memory_usage; }

    void prepare_ann_range_search(const doris::VectorSearchUserParams& params);

    Status evaluate_ann_range_search(

            const std::vector<std::unique_ptr<segment_v2::IndexIterator>>& cid_to_index_iterators,

            const std::vector<ColumnId>& idx_to_cid,

            const std::vector<std::unique_ptr<segment_v2::ColumnIterator>>& column_iterators,

            const std::unordered_map<VExprContext*, std::unordered_map<ColumnId, VExpr*>>&

                    common_expr_to_slotref_map,

            roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats,

            bool enable_result_cache, bool* ann_range_search_executed);

    uint64_t get_digest(uint64_t seed) const;

private:

    // Close method is called in vexpr context dector, not need call expicility

    void close();

    static void _reset_memory_usage(const VExprContextSPtrs& contexts);

    friend class VExpr;

    /// The expr tree this context is for.

    VExprSPtr _root;

    /// True if this context came from a Clone() call. Used to manage FunctionStateScope.

    bool _is_clone = false;

    /// Variables keeping track of current state.

    bool _prepared = false;

    bool _opened = false;

    /// FunctionContexts for each registered expression. The FunctionContexts are created

    /// and owned by this VExprContext.

    std::vector<std::unique_ptr<FunctionContext>> _fn_contexts;

    int _last_result_column_id = -1;

    /// The depth of expression-tree.

    int _depth_num = 0;

    std::shared_ptr<IndexExecContext> _index_context;

    ScanFilterHandle _scan_filter_handle;

    size_t _memory_usage = 0;

    segment_v2::AnnRangeSearchRuntime _ann_range_search_runtime;

    bool _suitable_for_ann_index = true;

    std::unique_ptr<RuntimeFilterSelectivity> _rf_selectivity =

            std::make_unique<RuntimeFilterSelectivity>();

};

} // namespace doris