// 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 "storage/segment/binary_dict_page.h"

#include <gen_cpp/olap_file.pb.h>

#include <gen_cpp/segment_v2.pb.h>

#include <algorithm>

#include <ostream>

#include <utility>

#include "common/compiler_util.h" // IWYU pragma: keep

#include "common/config.h"

#include "common/logging.h"

#include "common/status.h"

#include "core/column/column.h"

#include "core/column/column_string.h"

#include "storage/segment/binary_plain_page_v2.h"

#include "storage/segment/bitshuffle_page.h"

#include "storage/segment/encoding_info.h"

#include "util/coding.h"

#include "util/slice.h" // for Slice

namespace doris {

struct StringRef;

namespace segment_v2 {

BinaryDictPageBuilder::BinaryDictPageBuilder(const PageBuilderOptions& options)

        : _options(options),

          _finished(false),

          _data_page_builder(nullptr),

          _dict_builder(nullptr),

          _encoding_type(DICT_ENCODING),

          _dict_word_page_encoding_type(

                  options.encoding_preference.binary_plain_encoding_default_impl ==

                                  BinaryPlainEncodingTypePB::BINARY_PLAIN_ENCODING_V2

                          ? PLAIN_ENCODING_V2

                          : PLAIN_ENCODING),

          _fallback_binary_encoding_type(

                  options.encoding_preference.binary_plain_encoding_default_impl ==

                                  BinaryPlainEncodingTypePB::BINARY_PLAIN_ENCODING_V2

                          ? PLAIN_ENCODING_V2

                          : PLAIN_ENCODING) {}

Status BinaryDictPageBuilder::init() {

    // initially use DICT_ENCODING

    // TODO: the data page builder type can be created by Factory according to user config

    PageBuilder* data_page_builder_ptr = nullptr;

    RETURN_IF_ERROR(BitshufflePageBuilder<FieldType::OLAP_FIELD_TYPE_INT>::create(

            &data_page_builder_ptr, _options));

    _data_page_builder.reset(data_page_builder_ptr);

    PageBuilderOptions dict_builder_options;

    // here the binary plain page is used to store the dictionary items so

    // the data page size is set to the same as the dict page size

    dict_builder_options.data_page_size = _options.dict_page_size;

    dict_builder_options.dict_page_size = _options.dict_page_size;

    dict_builder_options.is_dict_page = true;

    const EncodingInfo* encoding_info;

    RETURN_IF_ERROR(EncodingInfo::get(FieldType::OLAP_FIELD_TYPE_VARCHAR,

                                      _dict_word_page_encoding_type, {}, &encoding_info));

    RETURN_IF_ERROR(encoding_info->create_page_builder(dict_builder_options, _dict_builder));

    return reset();

}

bool BinaryDictPageBuilder::is_page_full() {

    if (_data_page_builder->is_page_full()) {

        return true;

    }

    if (_encoding_type == DICT_ENCODING && _dict_builder->is_page_full()) {

        return true;

    }

    return false;

}

Status BinaryDictPageBuilder::add(const uint8_t* vals, size_t* count) {

    if (_encoding_type == DICT_ENCODING) {

        DCHECK(!_finished);

        DCHECK_GT(*count, 0);

        const Slice* src = reinterpret_cast<const Slice*>(vals);

        size_t num_added = 0;

        uint32_t value_code = -1;

        auto* actual_builder = dynamic_cast<BitshufflePageBuilder<FieldType::OLAP_FIELD_TYPE_INT>*>(

                _data_page_builder.get());

        for (int i = 0; i < *count; ++i, ++src) {

            if (is_page_full()) {

                break;

            }

            if (src->empty() && _has_empty) {

                value_code = _empty_code;

            } else if (auto iter = _dictionary.find(*src); iter != _dictionary.end()) {

                value_code = iter->second;

            } else {

                Slice dict_item(src->data, src->size);

                if (src->size > 0) {

                    char* item_mem = _arena.alloc(src->size);

                    if (item_mem == nullptr) {

                        return Status::MemoryAllocFailed("memory allocate failed, size:{}",

                                                         src->size);

                    }

                    dict_item.relocate(item_mem);

                }

                value_code = cast_set<uint32_t>(_dictionary.size());

                size_t add_count = 1;

                RETURN_IF_ERROR(_dict_builder->add(reinterpret_cast<const uint8_t*>(&dict_item),

                                                   &add_count));

                if (add_count == 0) {

                    // current dict page is full, stop processing remaining inputs

                    break;

                }

                _dictionary.emplace(dict_item, value_code);

                if (src->empty()) {

                    _has_empty = true;

                    _empty_code = value_code;

                }

            }

            size_t add_count = 1;

            RETURN_IF_ERROR(actual_builder->single_add(

                    reinterpret_cast<const uint8_t*>(&value_code), &add_count));

            if (add_count == 0) {

                // current data page is full, stop processing remaining inputs

                break;

            }

            // Track raw data size: the original string size

            _raw_data_size += src->size;

            num_added += 1;

        }

        *count = num_added;

        return Status::OK();

    } else {

        DCHECK(_encoding_type == PLAIN_ENCODING || _encoding_type == PLAIN_ENCODING_V2);

        RETURN_IF_ERROR(_data_page_builder->add(vals, count));

        // For plain encoding, track raw data size from the input

        const Slice* src = reinterpret_cast<const Slice*>(vals);

        for (size_t i = 0; i < *count; ++i) {

            _raw_data_size += src[i].size;

        }

        return Status::OK();

    }

}

Status BinaryDictPageBuilder::finish(OwnedSlice* slice) {

    if (VLOG_DEBUG_IS_ON && _encoding_type == DICT_ENCODING) {

        VLOG_DEBUG << "dict page size:" << _dict_builder->size();

    }

    DCHECK(!_finished);

    _finished = true;

    OwnedSlice data_slice;

    RETURN_IF_ERROR(_data_page_builder->finish(&data_slice));

    // TODO(gaodayue) separate page header and content to avoid this copy

    RETURN_IF_CATCH_EXCEPTION(

            { _buffer.append(data_slice.slice().data, data_slice.slice().size); });

    encode_fixed32_le(&_buffer[0], _encoding_type);

    *slice = _buffer.build();

    return Status::OK();

}

Status BinaryDictPageBuilder::reset() {

    RETURN_IF_CATCH_EXCEPTION({

        _finished = false;

        _raw_data_size = 0;

        _buffer.reserve(_options.data_page_size + BINARY_DICT_PAGE_HEADER_SIZE);

        _buffer.resize(BINARY_DICT_PAGE_HEADER_SIZE);

        if (_encoding_type == DICT_ENCODING && _dict_builder->is_page_full()) {

            DCHECK(_fallback_binary_encoding_type == PLAIN_ENCODING ||

                   _fallback_binary_encoding_type == PLAIN_ENCODING_V2);

            const EncodingInfo* encoding_info;

            RETURN_IF_ERROR(EncodingInfo::get(FieldType::OLAP_FIELD_TYPE_VARCHAR,

                                              _fallback_binary_encoding_type, {}, &encoding_info));

            RETURN_IF_ERROR(encoding_info->create_page_builder(_options, _data_page_builder));

            _encoding_type = _fallback_binary_encoding_type;

        } else {

            RETURN_IF_ERROR(_data_page_builder->reset());

        }

    });

    return Status::OK();

}

size_t BinaryDictPageBuilder::count() const {

    return _data_page_builder->count();

}

uint64_t BinaryDictPageBuilder::size() const {

    return _arena.used_size() + _data_page_builder->size();

}

Status BinaryDictPageBuilder::get_dictionary_page(OwnedSlice* dictionary_page) {

    return _dict_builder->finish(dictionary_page);

}

Status BinaryDictPageBuilder::get_dictionary_page_encoding(EncodingTypePB* encoding) const {

    *encoding = _dict_word_page_encoding_type;

    return Status::OK();

}

uint64_t BinaryDictPageBuilder::get_raw_data_size() const {

    return _raw_data_size;

}

BinaryDictPageDecoder::BinaryDictPageDecoder(Slice data, const PageDecoderOptions& options)

        : _data(data),

          _options(options),

          _data_page_decoder(nullptr),

          _parsed(false),

          _encoding_type(UNKNOWN_ENCODING) {}

Status BinaryDictPageDecoder::init() {

    CHECK(!_parsed);

    if (_data.size < BINARY_DICT_PAGE_HEADER_SIZE) {

        return Status::Corruption("invalid data size:{}, header size:{}", _data.size,

                                  BINARY_DICT_PAGE_HEADER_SIZE);

    }

    size_t type = decode_fixed32_le((const uint8_t*)&_data.data[0]);

    _encoding_type = static_cast<EncodingTypePB>(type);

    _data.remove_prefix(BINARY_DICT_PAGE_HEADER_SIZE);

    if (_encoding_type == DICT_ENCODING) {

        _data_page_decoder.reset(

                _bit_shuffle_ptr =

                        new BitShufflePageDecoder<FieldType::OLAP_FIELD_TYPE_INT>(_data, _options));

    } else if (_encoding_type == PLAIN_ENCODING) {

        _data_page_decoder.reset(

                new BinaryPlainPageDecoder<FieldType::OLAP_FIELD_TYPE_VARCHAR>(_data, _options));

    } else if (_encoding_type == PLAIN_ENCODING_V2) {

        _data_page_decoder.reset(

                new BinaryPlainPageV2Decoder<FieldType::OLAP_FIELD_TYPE_VARCHAR>(_data, _options));

    } else {

        LOG(WARNING) << "invalid encoding type:" << _encoding_type;

        return Status::Corruption("invalid encoding type:{}", _encoding_type);

    }

    RETURN_IF_ERROR(_data_page_decoder->init());

    _parsed = true;

    return Status::OK();

}

BinaryDictPageDecoder::~BinaryDictPageDecoder() {}

Status BinaryDictPageDecoder::seek_to_position_in_page(size_t pos) {

    return _data_page_decoder->seek_to_position_in_page(pos);

}

bool BinaryDictPageDecoder::is_dict_encoding() const {

    return _encoding_type == DICT_ENCODING;

}

void BinaryDictPageDecoder::set_dict_decoder(uint32_t num_dict_items, StringRef* dict_word_info) {

    _num_dict_items = num_dict_items;

    _dict_word_info = dict_word_info;

};

Status BinaryDictPageDecoder::next_batch(size_t* n, MutableColumnPtr& dst) {

    if (!is_dict_encoding()) {

        dst = dst->convert_to_predicate_column_if_dictionary();

        return _data_page_decoder->next_batch(n, dst);

    }

    // dictionary encoding

    DCHECK(_parsed);

    DCHECK(_dict_word_info != nullptr) << "_dict_word_info is nullptr";

    if (*n == 0 || _bit_shuffle_ptr->_cur_index >= _bit_shuffle_ptr->_num_elements) [[unlikely]] {

        *n = 0;

        return Status::OK();

    }

    size_t max_fetch = std::min(*n, static_cast<size_t>(_bit_shuffle_ptr->_num_elements -

                                                        _bit_shuffle_ptr->_cur_index));

    *n = max_fetch;

    if (_options.only_read_offsets) {

        // OFFSET_ONLY mode: resolve dict codes to get real string lengths

        // without copying actual char data. This allows length() to work.

        const auto* data_array = reinterpret_cast<const int32_t*>(_bit_shuffle_ptr->get_data(0));

        size_t start_index = _bit_shuffle_ptr->_cur_index;

        // Reuse _buffer (int32_t vector) to store uint32_t lengths.

        // int32_t and uint32_t have the same size/alignment, and string

        // lengths are always non-negative, so the bit patterns are identical.

        _buffer.resize(max_fetch);

        for (size_t i = 0; i < max_fetch; ++i) {

            int32_t codeword = data_array[start_index + i];

            _buffer[i] = static_cast<int32_t>(_dict_word_info[codeword].size);

        }

        dst->insert_offsets_from_lengths(reinterpret_cast<const uint32_t*>(_buffer.data()),

                                         max_fetch);

    } else {

        const auto* data_array = reinterpret_cast<const int32_t*>(_bit_shuffle_ptr->get_data(0));

        size_t start_index = _bit_shuffle_ptr->_cur_index;

        dst->insert_many_dict_data(data_array, start_index, _dict_word_info, max_fetch,

                                   _num_dict_items);

    }

    _bit_shuffle_ptr->_cur_index += max_fetch;

    return Status::OK();

}

Status BinaryDictPageDecoder::read_by_rowids(const rowid_t* rowids, ordinal_t page_first_ordinal,

                                             size_t* n, MutableColumnPtr& dst) {

    if (!is_dict_encoding()) {

        dst = dst->convert_to_predicate_column_if_dictionary();

        return _data_page_decoder->read_by_rowids(rowids, page_first_ordinal, n, dst);

    }

    DCHECK(_parsed);

    DCHECK(_dict_word_info != nullptr) << "_dict_word_info is nullptr";

    if (*n == 0) [[unlikely]] {

        *n = 0;

        return Status::OK();

    }

    auto total = *n;

    if (_options.only_read_offsets) {

        // OFFSET_ONLY mode: resolve dict codes to get real string lengths

        // without copying actual char data. This allows length() to work correctly.

        const auto* data_array = reinterpret_cast<const int32_t*>(_bit_shuffle_ptr->get_data(0));

        size_t read_count = 0;

        _buffer.resize(total);

        for (size_t i = 0; i < total; ++i) {

            ordinal_t ord = rowids[i] - page_first_ordinal;

            if (ord >= _bit_shuffle_ptr->_num_elements) [[unlikely]] {

                break;

            }

            int32_t codeword = data_array[ord];

            _buffer[read_count] = static_cast<int32_t>(_dict_word_info[codeword].size);

            read_count++;

        }

        if (read_count > 0) {

            dst->insert_offsets_from_lengths(reinterpret_cast<const uint32_t*>(_buffer.data()),

                                             read_count);

        }

        *n = read_count;

        return Status::OK();

    }

    const auto* data_array = reinterpret_cast<const int32_t*>(_bit_shuffle_ptr->get_data(0));

    size_t read_count = 0;

    _buffer.resize(total);

    for (size_t i = 0; i < total; ++i) {

        ordinal_t ord = rowids[i] - page_first_ordinal;

        if (ord >= _bit_shuffle_ptr->_num_elements) [[unlikely]] {

            break;

        }

        _buffer[read_count++] = data_array[ord];

    }

    if (LIKELY(read_count > 0)) {

        dst->insert_many_dict_data(_buffer.data(), 0, _dict_word_info, read_count, _num_dict_items);

    }

    *n = read_count;

    return Status::OK();

}

} // namespace segment_v2

} // namespace doris