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

#include <butil/macros.h>

#include <butil/time.h>

#include <gen_cpp/Metrics_types.h>

#include <parallel_hashmap/phmap.h>

#include <stdint.h>

#include <string.h>

#include <algorithm>

#include <memory>

#include <mutex>

#include <optional>

#include <ostream>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include <utility>

#include <vector>

#include "butil/containers/doubly_buffered_data.h"

#include "common/config.h"

#include "common/logging.h"

#include "common/status.h"

#include "core/block/block.h"

#include "core/data_type_serde/data_type_serde.h"

#include "exprs/vexpr_fwd.h"

#include "io/cache/remote_scan_cache_write_limiter.h"

#include "runtime/descriptors.h"

#include "runtime/exec_env.h"

#include "runtime/runtime_profile.h"

#include "storage/olap_common.h"

#include "storage/rowset/rowset.h"

#include "storage/tablet/tablet.h"

#include "storage/utils.h"

#include "util/lru_cache.h"

#include "util/mysql_global.h"

#include "util/slice.h"

namespace doris {

class PTabletKeyLookupRequest;

class PTabletKeyLookupResponse;

class RuntimeState;

class TDescriptorTable;

class TExpr;

// For caching point lookup pre allocted blocks and exprs

class Reusable {

public:

    ~Reusable();

    bool is_expired(int64_t ttl_ms) const {

        return butil::gettimeofday_ms() - _create_timestamp > ttl_ms;

    }

    Status init(const TDescriptorTable& t_desc_tbl, const std::vector<TExpr>& output_exprs,

                const TQueryOptions& query_options, const TabletSchema& schema,

                size_t block_size = 1);

    std::unique_ptr<Block> get_block();

    const DataTypeSerDeSPtrs& get_data_type_serdes() const { return _data_type_serdes; }

    const std::unordered_map<uint32_t, uint32_t>& get_col_uid_to_idx() const {

        return _col_uid_to_idx;

    }

    const std::vector<std::string>& get_col_default_values() const { return _col_default_values; }

    // do not touch block after returned

    void return_block(std::unique_ptr<Block>& block);

    TupleDescriptor* tuple_desc() { return _desc_tbl->get_tuple_descriptor(0); }

    const VExprContextSPtrs& output_exprs() { return _output_exprs_ctxs; }

    int32_t rs_column_uid() const { return _row_store_column_ids; }

    const std::unordered_set<int32_t> missing_col_uids() const { return _missing_col_uids; }

    const std::unordered_set<int32_t> include_col_uids() const { return _include_col_uids; }

    const std::unordered_set<int32_t>& column_store_col_uids() const {

        return _column_store_col_uids;

    }

    // Whether the projection holds VERSION/COMMIT_TSO, whose value depends on the owning rowset.

    bool has_rowset_derived_hidden_columns() const { return _has_rowset_derived_hidden_columns; }

    // Whether the row-store JSONB still serves any projected column once the rowset-derived

    // hidden columns are left out of it.

    bool decode_row_store() const { return _row_store_column_ids != -1 && _decode_row_store; }

    RuntimeState* runtime_state() { return _runtime_state.get(); }

    // delete sign idx in block

    int32_t delete_sign_idx() const { return _delete_sign_idx; }

private:

    // caching TupleDescriptor, output_expr, etc...

    std::unique_ptr<RuntimeState> _runtime_state;

    DescriptorTbl* _desc_tbl = nullptr;

    std::mutex _block_mutex;

    // prevent from allocte too many tmp blocks

    std::vector<std::unique_ptr<Block>> _block_pool;

    VExprContextSPtrs _output_exprs_ctxs;

    int64_t _create_timestamp = 0;

    DataTypeSerDeSPtrs _data_type_serdes;

    std::unordered_map<uint32_t, uint32_t> _col_uid_to_idx;

    std::vector<std::string> _col_default_values;

    // picked rowstore(column group) column unique id

    int32_t _row_store_column_ids = -1;

    // some column is missing in rowstore(column group), we need to fill them with column store values

    std::unordered_set<int32_t> _missing_col_uids;

    // cids decoded from the rowstore(column group) JSONB; empty means every slot

    std::unordered_set<int32_t> _include_col_uids;

    // Missing columns plus the projected VERSION/COMMIT_TSO: everything the JSONB does not serve.

    std::unordered_set<int32_t> _column_store_col_uids;

    bool _has_rowset_derived_hidden_columns = false;

    bool _decode_row_store = true;

    // delete sign idx in block

    int32_t _delete_sign_idx = -1;

};

// RowCache is a LRU cache for row store

class RowCache : public LRUCachePolicy {

public:

    using LRUCachePolicy::insert;

    // The cache key for row lru cache

    struct RowCacheKey {

        RowCacheKey(int64_t tablet_id, const Slice& key) : tablet_id(tablet_id), key(key) {}

        int64_t tablet_id;

        Slice key;

        // Encode to a flat binary which can be used as LRUCache's key

        std::string encode() const {

            std::string full_key;

            full_key.reserve(sizeof(int64_t) + key.size);

            const char* tid = reinterpret_cast<const char*>(&tablet_id);

            full_key.append(tid, tid + sizeof(int64_t));

            full_key.append(key.data, key.size);

            return full_key;

        }

    };

    class RowCacheValue : public LRUCacheValueBase {

    public:

        ~RowCacheValue() override { free(cache_value); }

        char* cache_value;

    };

    // A handle for RowCache entry. This class make it easy to handle

    // Cache entry. Users don't need to release the obtained cache entry. This

    // class will release the cache entry when it is destroyed.

    class CacheHandle {

    public:

        CacheHandle() = default;

        CacheHandle(LRUCachePolicy* cache, Cache::Handle* handle)

                : _cache(cache), _handle(handle) {}

        ~CacheHandle() {

            if (_handle != nullptr) {

                _cache->release(_handle);

            }

        }

        CacheHandle(CacheHandle&& other) noexcept {

            std::swap(_cache, other._cache);

            std::swap(_handle, other._handle);

        }

        CacheHandle& operator=(CacheHandle&& other) noexcept {

            std::swap(_cache, other._cache);

            std::swap(_handle, other._handle);

            return *this;

        }

        bool valid() { return _cache != nullptr && _handle != nullptr; }

        LRUCachePolicy* cache() const { return _cache; }

        Slice data() const {

            return {((RowCacheValue*)_cache->value(_handle))->cache_value,

                    reinterpret_cast<LRUHandle*>(_handle)->charge};

        }

    private:

        LRUCachePolicy* _cache = nullptr;

        Cache::Handle* _handle = nullptr;

        // Don't allow copy and assign

        DISALLOW_COPY_AND_ASSIGN(CacheHandle);

    };

    // Create global instance of this class

    static RowCache* create_global_cache(int64_t capacity, uint32_t num_shards = kDefaultNumShards);

    static RowCache* instance();

    // Lookup a row key from cache,

    // If the Row key is found, the cache entry will be written into handle.

    // CacheHandle will release cache entry to cache when it destructs

    // Return true if entry is found, otherwise return false.

    bool lookup(const RowCacheKey& key, CacheHandle* handle);

    // Insert a row with key into this cache.

    // This function is thread-safe, and when two clients insert two same key

    // concurrently, this function can assure that only one page is cached.

    // The in_memory page will have higher priority.

    void insert(const RowCacheKey& key, const Slice& data);

    //

    void erase(const RowCacheKey& key);

private:

    static constexpr uint32_t kDefaultNumShards = 128;

    RowCache(int64_t capacity, int num_shards = kDefaultNumShards);

};

// A cache used for prepare stmt.

// One connection per stmt perf uuid

class LookupConnectionCache : public LRUCachePolicy {

public:

    static LookupConnectionCache* instance() {

        return ExecEnv::GetInstance()->get_lookup_connection_cache();

    }

    static LookupConnectionCache* create_global_instance(size_t capacity);

private:

    friend class PointQueryExecutor;

    LookupConnectionCache(size_t capacity)

            : LRUCachePolicy(CachePolicy::CacheType::LOOKUP_CONNECTION_CACHE, capacity,

                             LRUCacheType::NUMBER, config::tablet_lookup_cache_stale_sweep_time_sec,

                             /*num shards*/ 32, /*element count capacity */ 0,

                             /*enable prune*/ true, /*is lru-k*/ true) {}

    static std::string encode_key(__int128_t cache_id) {

        fmt::memory_buffer buffer;

        fmt::format_to(buffer, "{}", cache_id);

        return std::string(buffer.data(), buffer.size());

    }

    void add(__int128_t cache_id, std::shared_ptr<Reusable> item) {

        std::string key = encode_key(cache_id);

        auto* value = new CacheValue;

        value->item = item;

        VLOG_DEBUG << "Add item mem"

                   << ", cache_capacity: " << get_capacity() << ", cache_usage: " << get_usage()

                   << ", mem_consum: " << mem_consumption();

        auto* lru_handle = insert(key, value, 1, sizeof(Reusable), CachePriority::NORMAL);

        release(lru_handle);

    }

    std::shared_ptr<Reusable> get(__int128_t cache_id) {

        std::string key = encode_key(cache_id);

        auto* lru_handle = lookup(key);

        if (lru_handle) {

            Defer release([cache = this, lru_handle] { cache->release(lru_handle); });

            auto* value = (CacheValue*)(LRUCachePolicy::value(lru_handle));

            return value->item;

        }

        return nullptr;

    }

    class CacheValue : public LRUCacheValueBase {

    public:

        ~CacheValue() override;

        std::shared_ptr<Reusable> item;

    };

};

struct Metrics {

    Metrics()

            : init_ns(TUnit::TIME_NS),

              init_key_ns(TUnit::TIME_NS),

              lookup_key_ns(TUnit::TIME_NS),

              lookup_data_ns(TUnit::TIME_NS),

              output_data_ns(TUnit::TIME_NS),

              load_segment_key_stage_ns(TUnit::TIME_NS),

              load_segment_data_stage_ns(TUnit::TIME_NS) {}

    RuntimeProfile::Counter init_ns;

    RuntimeProfile::Counter init_key_ns;

    RuntimeProfile::Counter lookup_key_ns;

    RuntimeProfile::Counter lookup_data_ns;

    RuntimeProfile::Counter output_data_ns;

    RuntimeProfile::Counter load_segment_key_stage_ns;

    RuntimeProfile::Counter load_segment_data_stage_ns;

    OlapReaderStatistics read_stats;

    size_t row_cache_hits = 0;

    bool hit_lookup_cache = false;

    size_t result_data_bytes;

};

// An util to do tablet lookup

class PointQueryExecutor {

public:

    ~PointQueryExecutor();

    Status init(const PTabletKeyLookupRequest* request, PTabletKeyLookupResponse* response);

    Status lookup_up();

    void print_profile();

    const OlapReaderStatistics& read_stats() const { return _read_stats; }

private:

    Status _init_keys(const PTabletKeyLookupRequest* request);

    Status _lookup_row_key();

    Status _lookup_row_data();

    Status _output_data();

    void _init_remote_scan_cache_write_limiter();

    static void release_rowset(RowsetSharedPtr* r) {

        if (r && *r) {

            VLOG_DEBUG << "release rowset " << (*r)->rowset_id();

            (*r)->release();

        }

        delete r;

    }

    // Read context for each row

    struct RowReadContext {

        RowReadContext() : _rowset_ptr(nullptr, &release_rowset) {}

        std::string _primary_key;

        RowCache::CacheHandle _cached_row_data;

        std::optional<RowLocation> _row_location;

        // rowset will be aquired during read

        // and released after used

        std::unique_ptr<RowsetSharedPtr, decltype(&release_rowset)> _rowset_ptr;

    };

    PTabletKeyLookupResponse* _response = nullptr;

    BaseTabletSPtr _tablet;

    std::vector<RowReadContext> _row_read_ctxs;

    std::shared_ptr<Reusable> _reusable;

    std::unique_ptr<Block> _result_block;

    Metrics _profile_metrics;

    bool _binary_row_format = false;

    OlapReaderStatistics _read_stats;

    std::unique_ptr<io::RemoteScanCacheWriteLimiter> _remote_scan_cache_write_limiter;

    int32_t _row_hits = 0;

    // snapshot read version

    int64_t _version = -1;

};

} // namespace doris