// 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 "format/table/paimon_cpp_reader.h"

#include <algorithm>

#include <mutex>

#include <tuple>

#include <unordered_map>

#include <utility>

#include "arrow/c/bridge.h"

#include "arrow/record_batch.h"

#include "arrow/result.h"

#include "core/block/block.h"

#include "core/block/column_with_type_and_name.h"

#include "format/table/paimon_doris_file_system.h"

#include "format/table/partition_column_filler.h"

#include "paimon/defs.h"

#include "paimon/memory/memory_pool.h"

#include "paimon/read_context.h"

#include "paimon/table/source/table_read.h"

#include "runtime/descriptors.h"

#include "runtime/runtime_state.h"

#include "util/url_coding.h"

namespace doris {

namespace {

constexpr const char* VALUE_KIND_FIELD = "_VALUE_KIND";

} // namespace

PaimonCppReader::PaimonCppReader(const std::vector<SlotDescriptor*>& file_slot_descs,

                                 RuntimeState* state, RuntimeProfile* profile,

                                 const TFileRangeDesc& range,

                                 const TFileScanRangeParams* range_params)

        : _file_slot_descs(file_slot_descs),

          _state(state),

          _profile(profile),

          _range(range),

          _range_params(range_params) {

    TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, _ctzz);

    if (range.__isset.table_format_params &&

        range.table_format_params.__isset.table_level_row_count) {

        _remaining_table_level_row_count = range.table_format_params.table_level_row_count;

    } else {

        _remaining_table_level_row_count = -1;

    }

}

PaimonCppReader::~PaimonCppReader() = default;

Status PaimonCppReader::on_before_init_reader(ReaderInitContext* ctx) {

    _column_descs = ctx->column_descs;

    _partition_values.clear();

    _partition_value_is_null.clear();

    if (ctx->range == nullptr || ctx->tuple_descriptor == nullptr ||

        !ctx->range->__isset.columns_from_path_keys) {

        return Status::OK();

    }

    DORIS_CHECK(ctx->range->__isset.columns_from_path);

    DORIS_CHECK(ctx->range->columns_from_path.size() == ctx->range->columns_from_path_keys.size());

    const bool has_null_flags = ctx->range->__isset.columns_from_path_is_null;

    if (has_null_flags) {

        DORIS_CHECK(ctx->range->columns_from_path_is_null.size() ==

                    ctx->range->columns_from_path_keys.size());

    }

    std::unordered_map<std::string, const SlotDescriptor*> name_to_slot;

    for (auto* slot : ctx->tuple_descriptor->slots()) {

        name_to_slot.emplace(slot->col_name(), slot);

    }

    for (size_t i = 0; i < ctx->range->columns_from_path_keys.size(); ++i) {

        const auto& key = ctx->range->columns_from_path_keys[i];

        auto slot_it = name_to_slot.find(key);

        if (slot_it == name_to_slot.end()) {

            continue;

        }

        _partition_values.emplace(

                key, std::make_tuple(ctx->range->columns_from_path[i], slot_it->second));

        _partition_value_is_null.emplace(

                key, has_null_flags ? ctx->range->columns_from_path_is_null[i] : false);

    }

    return Status::OK();

}

Status PaimonCppReader::on_after_read_block(Block* block, size_t* read_rows) {

    if (_column_descs == nullptr || _partition_values.empty() || *read_rows == 0 ||

        _push_down_agg_type == TPushAggOp::type::COUNT) {

        return Status::OK();

    }

    return _fill_partition_columns(block, *read_rows);

}

Status PaimonCppReader::init_reader() {

    if (_push_down_agg_type == TPushAggOp::type::COUNT && _remaining_table_level_row_count >= 0) {

        return Status::OK();

    }

    return _init_paimon_reader();

}

Status PaimonCppReader::_do_get_next_block(Block* block, size_t* read_rows, bool* eof) {

    if (_push_down_agg_type == TPushAggOp::type::COUNT && _remaining_table_level_row_count >= 0) {

        auto rows = std::min(_remaining_table_level_row_count,

                             (int64_t)_state->query_options().batch_size);

        _remaining_table_level_row_count -= rows;

        auto mutable_columns_guard = block->mutate_columns_scoped();

        auto& mutate_columns = mutable_columns_guard.mutable_columns();

        for (auto& col : mutate_columns) {

            col->resize(rows);

        }

        *read_rows = rows;

        *eof = false;

        if (_remaining_table_level_row_count == 0) {

            *eof = true;

        }

        return Status::OK();

    }

    if (!_batch_reader) {

        return Status::InternalError("paimon-cpp reader is not initialized");

    }

    if (_col_name_to_block_idx.empty()) {

        _col_name_to_block_idx = block->get_name_to_pos_map();

    }

    auto batch_result = _batch_reader->NextBatch();

    if (!batch_result.ok()) {

        return Status::InternalError("paimon-cpp read batch failed: {}",

                                     batch_result.status().ToString());

    }

    auto batch = std::move(batch_result).value();

    if (paimon::BatchReader::IsEofBatch(batch)) {

        *read_rows = 0;

        *eof = true;

        return Status::OK();

    }

    arrow::Result<std::shared_ptr<arrow::RecordBatch>> import_result =

            arrow::ImportRecordBatch(batch.first.get(), batch.second.get());

    if (!import_result.ok()) {

        return Status::InternalError("failed to import paimon-cpp arrow batch: {}",

                                     import_result.status().message());

    }

    auto record_batch = std::move(import_result).ValueUnsafe();

    const auto num_rows = static_cast<size_t>(record_batch->num_rows());

    const auto num_columns = record_batch->num_columns();

    auto columns_guard = block->mutate_columns_scoped();

    auto& columns = columns_guard.mutable_columns();

    for (int c = 0; c < num_columns; ++c) {

        const auto& field = record_batch->schema()->field(c);

        if (field->name() == VALUE_KIND_FIELD) {

            continue;

        }

        auto it = _col_name_to_block_idx.find(field->name());

        if (it == _col_name_to_block_idx.end()) {

            // Skip columns that are not in the block (e.g., partition columns handled elsewhere)

            continue;

        }

        const auto block_pos = it->second;

        try {

            RETURN_IF_ERROR(columns_guard.get_datatype_by_position(block_pos)

                                    ->get_serde()

                                    ->read_column_from_arrow(*columns[block_pos],

                                                             record_batch->column(c).get(), 0,

                                                             num_rows, _ctzz));

        } catch (Exception& e) {

            return Status::InternalError("Failed to convert from arrow to block: {}", e.what());

        }

    }

    *read_rows = num_rows;

    *eof = false;

    return Status::OK();

}

Status PaimonCppReader::_get_columns_impl(

        std::unordered_map<std::string, DataTypePtr>* name_to_type) {

    for (const auto& slot : _file_slot_descs) {

        name_to_type->emplace(slot->col_name(), slot->type());

    }

    return Status::OK();

}

Status PaimonCppReader::_fill_partition_columns(Block* block, size_t num_rows) {

    if (_col_name_to_block_idx.empty()) {

        _col_name_to_block_idx = block->get_name_to_pos_map();

    }

    for (const auto& desc : *_column_descs) {

        if (desc.category != ColumnCategory::PARTITION_KEY) {

            continue;

        }

        auto value_it = _partition_values.find(desc.name);

        if (value_it == _partition_values.end()) {

            continue;

        }

        auto col_it = _col_name_to_block_idx.find(desc.name);

        if (col_it == _col_name_to_block_idx.end()) {

            return Status::InternalError("Missing partition column {} in block {}", desc.name,

                                         block->dump_structure());

        }

        auto& column_with_type_and_name = block->get_by_position(col_it->second);

        auto mutable_column = std::move(*column_with_type_and_name.column).mutate();

        const auto& [value, slot_desc] = value_it->second;

        auto null_it = _partition_value_is_null.find(desc.name);

        DORIS_CHECK(null_it != _partition_value_is_null.end());

        RETURN_IF_ERROR(fill_partition_column_from_path_value(*mutable_column, *slot_desc, value,

                                                              num_rows, null_it->second));

        column_with_type_and_name.column = std::move(mutable_column);

    }

    return Status::OK();

}

Status PaimonCppReader::close() {

    if (_batch_reader) {

        _batch_reader->Close();

    }

    return Status::OK();

}

Status PaimonCppReader::_init_paimon_reader() {

    register_paimon_doris_file_system();

    RETURN_IF_ERROR(_decode_split(&_split));

    auto table_path_opt = _resolve_table_path();

    if (!table_path_opt.has_value()) {

        return Status::InternalError(

                "paimon-cpp missing paimon_table; cannot resolve paimon table root path");

    }

    auto options = _build_options();

    auto read_columns = _build_read_columns();

    // Avoid moving strings across module boundaries to prevent allocator mismatches in ASAN builds.

    std::string table_path = table_path_opt.value();

    static std::once_flag options_log_once;

    std::call_once(options_log_once, [&]() {

        auto has_key = [&](const char* key) {

            auto it = options.find(key);

            return (it != options.end() && !it->second.empty()) ? "set" : "empty";

        };

        auto value_or = [&](const char* key) {

            auto it = options.find(key);

            return it != options.end() ? it->second : std::string("<unset>");

        };

        LOG(INFO) << "paimon-cpp options summary: table_path=" << table_path

                  << " AWS_ACCESS_KEY=" << has_key("AWS_ACCESS_KEY")

                  << " AWS_SECRET_KEY=" << has_key("AWS_SECRET_KEY")

                  << " AWS_TOKEN=" << has_key("AWS_TOKEN")

                  << " AWS_ENDPOINT=" << value_or("AWS_ENDPOINT")

                  << " AWS_REGION=" << value_or("AWS_REGION")

                  << " use_path_style=" << value_or("use_path_style")

                  << " fs.oss.endpoint=" << value_or("fs.oss.endpoint")

                  << " fs.s3a.endpoint=" << value_or("fs.s3a.endpoint");

    });

    paimon::ReadContextBuilder builder(table_path);

    if (!read_columns.empty()) {

        builder.SetReadSchema(read_columns);

    }

    if (!options.empty()) {

        builder.SetOptions(options);

    }

    if (_predicate) {

        builder.SetPredicate(_predicate);

        builder.EnablePredicateFilter(true);

    }

    auto context_result = builder.Finish();

    if (!context_result.ok()) {

        return Status::InternalError("paimon-cpp build read context failed: {}",

                                     context_result.status().ToString());

    }

    auto context = std::move(context_result).value();

    auto table_read_result = paimon::TableRead::Create(std::move(context));

    if (!table_read_result.ok()) {

        return Status::InternalError("paimon-cpp create table read failed: {}",

                                     table_read_result.status().ToString());

    }

    auto table_read = std::move(table_read_result).value();

    auto reader_result = table_read->CreateReader(_split);

    if (!reader_result.ok()) {

        return Status::InternalError("paimon-cpp create reader failed: {}",

                                     reader_result.status().ToString());

    }

    _table_read = std::move(table_read);

    _batch_reader = std::move(reader_result).value();

    return Status::OK();

}

Status PaimonCppReader::_decode_split(std::shared_ptr<paimon::Split>* split) {

    if (!_range.__isset.table_format_params || !_range.table_format_params.__isset.paimon_params ||

        !_range.table_format_params.paimon_params.__isset.paimon_split) {

        return Status::InternalError("paimon-cpp missing paimon_split in scan range");

    }

    const auto& encoded_split = _range.table_format_params.paimon_params.paimon_split;

    std::string decoded_split;

    if (!base64_decode(encoded_split, &decoded_split)) {

        return Status::InternalError("paimon-cpp base64 decode paimon_split failed");

    }

    auto pool = paimon::GetDefaultPool();

    auto split_result =

            paimon::Split::Deserialize(decoded_split.data(), decoded_split.size(), pool);

    if (!split_result.ok()) {

        return Status::InternalError("paimon-cpp deserialize split failed: {}",

                                     split_result.status().ToString());

    }

    *split = std::move(split_result).value();

    return Status::OK();

}

std::optional<std::string> PaimonCppReader::_resolve_table_path() const {

    if (_range.__isset.table_format_params && _range.table_format_params.__isset.paimon_params &&

        _range.table_format_params.paimon_params.__isset.paimon_table &&

        !_range.table_format_params.paimon_params.paimon_table.empty()) {

        return _range.table_format_params.paimon_params.paimon_table;

    }

    return std::nullopt;

}

std::vector<std::string> PaimonCppReader::_build_read_columns() const {

    std::vector<std::string> columns;

    columns.reserve(_file_slot_descs.size());

    for (const auto& slot : _file_slot_descs) {

        columns.emplace_back(slot->col_name());

    }

    return columns;

}

std::map<std::string, std::string> PaimonCppReader::_build_options() const {

    std::map<std::string, std::string> options;

    if (_range_params && _range_params->__isset.paimon_options &&

        !_range_params->paimon_options.empty()) {

        options.insert(_range_params->paimon_options.begin(), _range_params->paimon_options.end());

    } else if (_range.__isset.table_format_params &&

               _range.table_format_params.__isset.paimon_params &&

               _range.table_format_params.paimon_params.__isset.paimon_options) {

        options.insert(_range.table_format_params.paimon_params.paimon_options.begin(),

                       _range.table_format_params.paimon_params.paimon_options.end());

    }

    if (_range_params && _range_params->__isset.properties && !_range_params->properties.empty()) {

        for (const auto& kv : _range_params->properties) {

            options[kv.first] = kv.second;

        }

    } else if (_range.__isset.table_format_params &&

               _range.table_format_params.__isset.paimon_params &&

               _range.table_format_params.paimon_params.__isset.hadoop_conf) {

        for (const auto& kv : _range.table_format_params.paimon_params.hadoop_conf) {

            options[kv.first] = kv.second;

        }

    }

    auto copy_if_missing = [&](const char* from_key, const char* to_key) {

        if (options.find(to_key) != options.end()) {

            return;

        }

        auto it = options.find(from_key);

        if (it != options.end() && !it->second.empty()) {

            options[to_key] = it->second;

        }

    };

    // Map common OSS/S3 Hadoop configs to Doris S3 property keys.

    copy_if_missing("fs.oss.accessKeyId", "AWS_ACCESS_KEY");

    copy_if_missing("fs.oss.accessKeySecret", "AWS_SECRET_KEY");

    copy_if_missing("fs.oss.sessionToken", "AWS_TOKEN");

    copy_if_missing("fs.oss.endpoint", "AWS_ENDPOINT");

    copy_if_missing("fs.oss.region", "AWS_REGION");

    copy_if_missing("fs.s3a.access.key", "AWS_ACCESS_KEY");

    copy_if_missing("fs.s3a.secret.key", "AWS_SECRET_KEY");

    copy_if_missing("fs.s3a.session.token", "AWS_TOKEN");

    copy_if_missing("fs.s3a.endpoint", "AWS_ENDPOINT");

    copy_if_missing("fs.s3a.region", "AWS_REGION");

    copy_if_missing("fs.s3a.path.style.access", "use_path_style");

    // Backfill file.format/manifest.format from split file_format to avoid

    // paimon-cpp falling back to default manifest.format=avro.

    if (_range.__isset.table_format_params && _range.table_format_params.__isset.paimon_params &&

        _range.table_format_params.paimon_params.__isset.file_format &&

        !_range.table_format_params.paimon_params.file_format.empty()) {

        const auto& split_file_format = _range.table_format_params.paimon_params.file_format;

        auto file_format_it = options.find(paimon::Options::FILE_FORMAT);

        if (file_format_it == options.end() || file_format_it->second.empty()) {

            options[paimon::Options::FILE_FORMAT] = split_file_format;

        }

        auto manifest_format_it = options.find(paimon::Options::MANIFEST_FORMAT);

        if (manifest_format_it == options.end() || manifest_format_it->second.empty()) {

            options[paimon::Options::MANIFEST_FORMAT] = split_file_format;

        }

    }

    options[paimon::Options::FILE_SYSTEM] = "doris";

    return options;

}

} // namespace doris