// 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 <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 "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; }

    /// 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);

    [[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);

    [[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);

    [[nodiscard]] static Status execute_conjuncts_and_filter_block(
            const VExprContextSPtrs& ctxs, Block* block, std::vector<uint32_t>& columns_to_filter,
            int column_to_keep);

    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);

    [[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;
        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;
        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);

    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;
    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