// 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/segment_prefetcher.h"

#include <algorithm>

#include <ranges>

#include "common/config.h"

#include "common/logging.h"

#include "storage/index/ordinal_page_index.h"

#include "storage/iterators.h"

#include "storage/segment/column_reader.h"

namespace doris::segment_v2 {

void SegmentPrefetcher::add_rowids(const rowid_t* rowids, uint32_t num) {

    if (ordinal_index == nullptr) {

        return;

    }

    const auto& ordinals = ordinal_index->_ordinals; // ordinals[i] = first ordinal of page i

    const auto& pages = ordinal_index->_pages;       // pages[i] = page pointer of page i

    const int num_pages = ordinal_index->_num_pages;

    for (uint32_t i = 0; i < num; ++i) {

        rowid_t rowid = rowids[i];

        if (_is_forward) {

            while (page_idx < num_pages - 1 && ordinals[page_idx + 1] <= rowid) {

                page_idx++;

            }

            const auto& page = pages[page_idx];

            size_t page_start_block = _offset_to_block_id(page.offset);

            size_t page_end_block = _offset_to_block_id(page.offset + page.size - 1);

            // If page spans two blocks, assign it to the next block (page_end_block)

            size_t block_id =

                    (page_start_block != page_end_block) ? page_end_block : page_start_block;

            if (block_id != last_block_id) {

                if (last_block_id != static_cast<size_t>(-1)) {

                    _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

                }

                last_block_id = block_id;

                current_block_first_rowid = rowid;

            }

        } else {

            // Backward reading: we need the last rowid in each block as the "first" rowid

            // (because when reading backwards, we encounter the largest rowid first)

            //

            // Strategy: iterate forward through bitmap, but for each block,

            // keep updating current_block_first_rowid to the latest (largest) rowid in that block

            while (page_idx < num_pages - 1 && ordinals[page_idx + 1] <= rowid) {

                page_idx++;

            }

            size_t block_id = _offset_to_block_id(pages[page_idx].offset);

            if (block_id != last_block_id) {

                if (last_block_id != static_cast<size_t>(-1)) {

                    _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

                }

                last_block_id = block_id;

            }

            current_block_first_rowid = rowid;

        }

    }

}

void SegmentPrefetcher::build_all_data_blocks() {

    if (ordinal_index == nullptr) {

        return;

    }

    reset_blocks();

    const auto& ordinals = ordinal_index->_ordinals; // ordinals[i] = first ordinal of page i

    const auto& pages = ordinal_index->_pages;       // pages[i] = page pointer of page i

    const int num_pages = ordinal_index->_num_pages;

    last_block_id = static_cast<size_t>(-1);

    current_block_first_rowid = 0;

    for (page_idx = 0; page_idx < num_pages; ++page_idx) {

        const auto& page = pages[page_idx];

        if (_is_forward) {

            size_t page_start_block = _offset_to_block_id(page.offset);

            size_t page_end_block = _offset_to_block_id(page.offset + page.size - 1);

            // If page spans two blocks, assign it to the next block (page_end_block)

            size_t block_id =

                    (page_start_block != page_end_block) ? page_end_block : page_start_block;

            if (block_id != last_block_id) {

                if (last_block_id != static_cast<size_t>(-1)) {

                    _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

                }

                last_block_id = block_id;

                current_block_first_rowid = static_cast<rowid_t>(ordinals[page_idx]);

            }

        } else {

            // Backward: use the last ordinal in each block as first_rowid

            size_t block_id = _offset_to_block_id(page.offset);

            if (block_id != last_block_id) {

                if (last_block_id != static_cast<size_t>(-1)) {

                    _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

                }

                last_block_id = block_id;

            }

            current_block_first_rowid = static_cast<rowid_t>(ordinals[page_idx]);

        }

    }

    // Add the last block

    if (last_block_id != static_cast<size_t>(-1)) {

        _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

    }

    // Reverse for backward reading

    if (!_is_forward && !_block_sequence.empty()) {

        std::ranges::reverse(_block_sequence);

    }

}

void SegmentPrefetcher::build_blocks_by_rowids(const roaring::Roaring& row_bitmap,

                                               const std::vector<SegmentPrefetcher*>& prefetchers) {

    for (auto* prefetcher : prefetchers) {

        prefetcher->begin_build_blocks_by_rowids();

    }

    int batch_size = config::segment_file_cache_consume_rowids_batch_size;

    std::vector<rowid_t> rowids(batch_size);

    roaring::api::roaring_uint32_iterator_t iter;

    roaring::api::roaring_init_iterator(&row_bitmap.roaring, &iter);

    uint32_t num = roaring::api::roaring_read_uint32_iterator(&iter, rowids.data(), batch_size);

    for (; num > 0;

         num = roaring::api::roaring_read_uint32_iterator(&iter, rowids.data(), batch_size)) {

        for (auto* prefetcher : prefetchers) {

            prefetcher->add_rowids(rowids.data(), num);

        }

    }

    for (auto* prefetcher : prefetchers) {

        prefetcher->finish_build_blocks_by_rowids();

    }

}

void SegmentPrefetcher::begin_build_blocks_by_rowids() {

    reset_blocks();

    page_idx = 0;

}

void SegmentPrefetcher::finish_build_blocks_by_rowids() {

    if (ordinal_index == nullptr) {

        return;

    }

    if (last_block_id != static_cast<size_t>(-1)) {

        _block_sequence.emplace_back(last_block_id, current_block_first_rowid);

    }

    if (!_is_forward && !_block_sequence.empty()) {

        std::ranges::reverse(_block_sequence);

    }

    LOG_IF(INFO, config::enable_segment_prefetch_verbose_log) << fmt::format(

            "[verbose] SegmentPrefetcher initialized with block count={}, is_forward={}, "

            "num_pages={}, path={}, blocks: (block_id, first_rowid)=[{}]",

            _block_sequence.size(), _is_forward, ordinal_index->_num_pages, _path,

            fmt::join(_block_sequence | std::views::transform([](const auto& b) {

                          return fmt::format("({}, {})", b.block_id, b.first_rowid);

                      }),

                      ","));

}

void SegmentPrefetcher::reset_blocks() {

    _block_sequence.clear();

    _current_block_index = 0;

    _prefetched_index = -1;

}

Status SegmentPrefetcher::init(std::shared_ptr<ColumnReader> column_reader,

                               const StorageReadOptions& read_options) {

    DCHECK(column_reader != nullptr);

    reset_blocks();

    _is_forward = !read_options.read_orderby_key_reverse;

    _path = column_reader->_file_reader->path().filename().native();

    RETURN_IF_ERROR(column_reader->get_ordinal_index_reader(ordinal_index, read_options.stats));

    return Status::OK();

}

bool SegmentPrefetcher::need_prefetch(rowid_t current_rowid, std::vector<BlockRange>* out_ranges) {

    DCHECK(out_ranges != nullptr);

    LOG_IF(INFO, config::enable_segment_prefetch_verbose_log)

            << fmt::format("[verbose] SegmentPrefetcher need_prefetch enter current_rowid={}, {}",

                           current_rowid, debug_string());

    if (_block_sequence.empty() ||

        _prefetched_index >= static_cast<int>(_block_sequence.size()) - 1) {

        return false;

    }

    LOG_IF(INFO, config::enable_segment_prefetch_verbose_log) << fmt::format(

            "[verbose] SegmentPrefetcher need_prefetch called with current_rowid={}, {}, "

            "block=(id={}, first_rowid={})",

            current_rowid, debug_string(), _block_sequence[_current_block_index].block_id,

            _block_sequence[_current_block_index].first_rowid);

    if (_is_forward) {

        while (_current_block_index + 1 < _block_sequence.size() &&

               _block_sequence[_current_block_index + 1].first_rowid <= current_rowid) {

            _current_block_index++;

        }

    } else {

        while (_current_block_index + 1 < _block_sequence.size() &&

               _block_sequence[_current_block_index + 1].first_rowid >= current_rowid) {

            _current_block_index++;

        }

    }

    out_ranges->clear();

    // for non-predicate column, some rowids in row_bitmap may be filtered out after vec evaluation of predicate columns,

    // so we should not prefetch for these rows

    _prefetched_index = std::max(_prefetched_index, _current_block_index - 1);

    while (_prefetched_index + 1 < _block_sequence.size() &&

           window_size() < _config.prefetch_window_size) {

        out_ranges->push_back(_block_id_to_range(_block_sequence[++_prefetched_index].block_id));

    }

    LOG_IF(INFO, config::enable_segment_prefetch_verbose_log) << fmt::format(

            "[verbose] SegmentPrefetcher need_prefetch after calc with current_rowid={}, {}, "

            "block=(id={}, first_rowid={})",

            current_rowid, debug_string(), _block_sequence[_current_block_index].block_id,

            _block_sequence[_current_block_index].first_rowid);

    bool triggered = !out_ranges->empty();

    if (triggered) {

        LOG_IF(INFO, config::enable_segment_prefetch_verbose_log) << fmt::format(

                "[verbose] SegmentPrefetcher prefetch triggered at rowid={}, {}, prefetch {} "

                "blocks: (offset, size)=[{}]",

                current_rowid, debug_string(), out_ranges->size(),

                fmt::join(*out_ranges | std::views::transform([](const auto& b) {

                    return fmt::format("({}, {})", b.offset, b.size);

                }),

                          ","));

    }

    return triggered;

}

} // namespace doris::segment_v2