/*
 * 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 "paimon/format/parquet/file_reader_wrapper.h"

#include <cassert>
#include <cstddef>

#include "arrow/record_batch.h"
#include "arrow/util/range.h"
#include "fmt/format.h"
#include "paimon/macros.h"
#include "parquet/arrow/reader.h"
#include "parquet/file_reader.h"
#include "parquet/metadata.h"

namespace paimon::parquet {

Result<std::unique_ptr<FileReaderWrapper>> FileReaderWrapper::Create(
    std::unique_ptr<::parquet::arrow::FileReader>&& file_reader) {
    if (file_reader == nullptr) {
        return Status::Invalid("file reader wrapper create failed. file reader is nullptr");
    }
    std::vector<std::pair<uint64_t, uint64_t>> all_row_group_ranges;
    auto meta_data = file_reader->parquet_reader()->metadata();
    // prepare [start_row_idx, end_row_idx) for all row groups
    uint64_t start_row_idx = 0;
    for (int32_t i = 0; i < meta_data->num_row_groups(); i++) {
        uint64_t end_row_idx = start_row_idx + meta_data->RowGroup(i)->num_rows();
        all_row_group_ranges.emplace_back(start_row_idx, end_row_idx);
        start_row_idx = end_row_idx;
    }
    uint64_t num_rows = file_reader->parquet_reader()->metadata()->num_rows();
    if (start_row_idx != num_rows) {
        assert(false);
        return Status::Invalid(
            fmt::format("unexpected error. row group ranges not match with num rows {}", num_rows));
    }
    std::vector<int32_t> row_groups_indices = arrow::internal::Iota(file_reader->num_row_groups());
    std::vector<int32_t> columns_indices =
        arrow::internal::Iota(file_reader->parquet_reader()->metadata()->num_columns());
    auto file_reader_wrapper = std::unique_ptr<FileReaderWrapper>(
        new FileReaderWrapper(std::move(file_reader), all_row_group_ranges, num_rows));
    PAIMON_RETURN_NOT_OK(file_reader_wrapper->PrepareForReadingLazy(
        std::set<int32_t>(row_groups_indices.begin(), row_groups_indices.end()), columns_indices));
    return file_reader_wrapper;
}

FileReaderWrapper::FileReaderWrapper(
    std::unique_ptr<::parquet::arrow::FileReader>&& file_reader,
    const std::vector<std::pair<uint64_t, uint64_t>>& all_row_group_ranges, uint64_t num_rows)
    : file_reader_(std::move(file_reader)),
      all_row_group_ranges_(all_row_group_ranges),
      num_rows_(num_rows) {}

Status FileReaderWrapper::SeekToRow(uint64_t row_number) {
    for (uint64_t i = 0; i < target_row_groups_.size(); i++) {
        if (row_number > target_row_groups_[i].first && row_number < target_row_groups_[i].second) {
            return Status::Invalid(fmt::format(
                "seek to row failed. row number {} should not be in the middle of readable range",
                row_number));
        }
        if (target_row_groups_[i].first >= row_number) {
            current_row_group_idx_ = i;
            next_row_to_read_ = target_row_groups_[i].first;
            std::vector<int32_t> target_row_group_indices;
            for (uint64_t j = i; j < target_row_groups_.size(); j++) {
                PAIMON_ASSIGN_OR_RAISE(int32_t row_group_id, GetRowGroupId(target_row_groups_[j]));
                target_row_group_indices.push_back(row_group_id);
            }
            PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
                target_row_group_indices, target_column_indices_, &batch_reader_));
            return Status::OK();
        }
    }
    next_row_to_read_ = num_rows_;
    current_row_group_idx_ = target_row_groups_.size();
    return Status::OK();
}

Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::Next() {
    if (PAIMON_UNLIKELY(!reader_initialized_)) {
        PAIMON_RETURN_NOT_OK(PrepareForReading(target_row_group_indices_, target_column_indices_));
    }
    std::shared_ptr<arrow::RecordBatch> record_batch;
    if (current_row_group_idx_ < target_row_groups_.size()) {
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(record_batch, batch_reader_->Next());
    }
    if (record_batch) {
        int64_t num_rows = record_batch->num_rows();
        previous_first_row_ = next_row_to_read_;
        if (next_row_to_read_ + num_rows < target_row_groups_[current_row_group_idx_].second) {
            next_row_to_read_ += num_rows;
        } else if (next_row_to_read_ + num_rows ==
                   target_row_groups_[current_row_group_idx_].second) {
            if (current_row_group_idx_ == target_row_groups_.size() - 1) {
                // current row group is the last.
                next_row_to_read_ = num_rows_;
            } else {
                current_row_group_idx_++;
                next_row_to_read_ = target_row_groups_[current_row_group_idx_].first;
            }
        } else {
            return Status::Invalid(fmt::format(
                "Next failed. Unexpected error, next row to read {} + num rows just read {} "
                "should always be within current row group range or exactly equals to current "
                "row group end {}",
                next_row_to_read_, num_rows, target_row_groups_[current_row_group_idx_].second));
        }
    } else {
        previous_first_row_ = next_row_to_read_;
    }
    return record_batch;
}

Result<std::vector<std::pair<uint64_t, uint64_t>>> FileReaderWrapper::GetRowGroupRanges(
    const std::set<int32_t>& row_group_indices) const {
    std::vector<std::pair<uint64_t, uint64_t>> row_group_ranges;
    for (auto row_group_index : row_group_indices) {
        if (static_cast<size_t>(row_group_index) >= all_row_group_ranges_.size()) {
            return Status::Invalid(fmt::format("row group index {} is out of bound {}",
                                               row_group_index, all_row_group_ranges_.size()));
        }
        row_group_ranges.push_back(all_row_group_ranges_[row_group_index]);
    }
    return row_group_ranges;
}

Status FileReaderWrapper::PrepareForReadingLazy(const std::set<int32_t>& target_row_group_indices,
                                                const std::vector<int32_t>& column_indices) {
    target_row_group_indices_ = target_row_group_indices;
    target_column_indices_ = column_indices;
    reader_initialized_ = false;
    return Status::OK();
}

Status FileReaderWrapper::PrepareForReading(const std::set<int32_t>& target_row_group_indices,
                                            const std::vector<int32_t>& column_indices) {
    std::vector<std::pair<uint64_t, uint64_t>> target_row_groups;
    PAIMON_ASSIGN_OR_RAISE(target_row_groups, GetRowGroupRanges(target_row_group_indices));
    std::unique_ptr<arrow::RecordBatchReader> batch_reader;
    PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
        std::vector<int32_t>(target_row_group_indices.begin(), target_row_group_indices.end()),
        column_indices, &batch_reader));
    target_row_groups_ = target_row_groups;
    target_column_indices_ = column_indices;
    batch_reader_ = std::move(batch_reader);
    if (target_row_groups_.empty()) {
        next_row_to_read_ = num_rows_;
    } else {
        next_row_to_read_ = target_row_groups_[0].first;
    }
    previous_first_row_ = std::numeric_limits<uint64_t>::max();
    current_row_group_idx_ = 0;
    reader_initialized_ = true;
    return Status::OK();
}

Result<std::set<int32_t>> FileReaderWrapper::FilterRowGroupsByReadRanges(
    const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges,
    const std::vector<int32_t>& src_row_groups) const {
    std::set<int32_t> target_row_groups;
    PAIMON_ASSIGN_OR_RAISE(std::set<int32_t> row_groups_to_read,
                           ReadRangesToRowGroupIds(read_ranges));
    for (const auto& row_group_id : src_row_groups) {
        if (row_groups_to_read.find(row_group_id) != row_groups_to_read.end()) {
            target_row_groups.emplace(row_group_id);
        }
    }
    return target_row_groups;
}

Result<std::set<int32_t>> FileReaderWrapper::ReadRangesToRowGroupIds(
    const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges) const {
    std::set<int32_t> selected_row_group_ids;
    for (const auto& read_range : read_ranges) {
        PAIMON_ASSIGN_OR_RAISE(int32_t row_group_id, GetRowGroupId(read_range));
        selected_row_group_ids.emplace(row_group_id);
    }
    return selected_row_group_ids;
}

Result<int32_t> FileReaderWrapper::GetRowGroupId(std::pair<uint64_t, uint64_t> target_range) const {
    for (size_t i = 0; i < all_row_group_ranges_.size(); i++) {
        if (all_row_group_ranges_[i] == target_range) {
            return i;
        }
    }
    return Status::Invalid(fmt::format(
        "not expected failure. target range bound '{},{}' not match with row group range bound",
        target_range.first, target_range.second));
}

}  // namespace paimon::parquet