/*
 * 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 "benchmark/benchmark_suite.h"

#include <atomic>
#include <cstdint>
#include <iostream>
#include <map>
#include <memory>
#include <mutex>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "arrow/api.h"
#include "arrow/c/bridge.h"
#include "arrow/io/api.h"
#include "benchmark/benchmark_helpers.h"
#include "benchmark/cli_option_parsing.h"
#include "paimon/api.h"
#include "paimon/catalog/catalog.h"
#include "paimon/testing/utils/testharness.h"

#if __has_include("parquet/arrow/reader.h")
#include "parquet/arrow/reader.h"
#include "parquet/file_reader.h"
#define PAIMON_BENCHMARK_HAS_PARQUET_READER 1
#else
#define PAIMON_BENCHMARK_HAS_PARQUET_READER 0
#endif

namespace paimon::benchmark {

namespace {

constexpr int64_t kSourceBatchMaxRows = 4096;
constexpr int32_t kRowToBatchThreadNumber = 3;

struct BenchmarkCliOptions {
    std::string source_data_file;
    std::string source_table_path;
    std::vector<std::string> pk_columns;
    std::vector<std::pair<std::string, std::string>> extra_options;
};

struct SourceDataSpec {
    std::string format;
    std::string path;
};

BenchmarkCliOptions& MutableBenchmarkCliOptions() {
    static BenchmarkCliOptions options;
    return options;
}

const BenchmarkCliOptions& GetBenchmarkCliOptions() {
    return MutableBenchmarkCliOptions();
}

Status ParsePaimonBenchmarkCliArgsImpl(int32_t* argc, char** argv) {
    auto& options = MutableBenchmarkCliOptions();
    options = BenchmarkCliOptions{};
    const int32_t parsed_argc = *argc;
    int32_t write_index = 1;
    for (int32_t arg_index = 1; arg_index < parsed_argc; ++arg_index) {
        const std::string arg(argv[arg_index]);

        PAIMON_ASSIGN_OR_RAISE(bool is_parsed,
                               paimon::benchmark::ParseStringOptionArg(
                                   parsed_argc, argv, arg, "--paimon_source_data_file", &arg_index,
                                   &options.source_data_file));
        if (is_parsed) {
            continue;
        }
        PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseStringOptionArg(
                                              parsed_argc, argv, arg, "--paimon_source_table_path",
                                              &arg_index, &options.source_table_path));
        if (is_parsed) {
            continue;
        }
        PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseCommaSeparatedOptionArg(
                                              parsed_argc, argv, arg, "--paimon_pk_columns",
                                              &arg_index, &options.pk_columns));
        if (is_parsed) {
            continue;
        }
        PAIMON_ASSIGN_OR_RAISE(is_parsed, paimon::benchmark::ParseDelimitedRepeatableOptionArg(
                                              parsed_argc, argv, arg, "--paimon_option", &arg_index,
                                              &options.extra_options));
        if (is_parsed) {
            continue;
        }

        argv[write_index++] = argv[arg_index];
    }

    *argc = write_index;
    argv[write_index] = nullptr;
    return Status::OK();
}

bool HasHelpFlagImpl(int32_t argc, char** argv) {
    for (int32_t arg_index = 1; arg_index < argc; ++arg_index) {
        const std::string arg(argv[arg_index]);
        if (arg == "-h" || arg == "--help" || arg == "--help=true") {
            return true;
        }
    }
    return false;
}

void PrintPaimonBenchmarkCliHelpImpl() {
    std::cout
        << "Paimon benchmark custom options:\n"
        << "  --paimon_source_data_file=<path>\n"
        << "      Required. External source data file used to build benchmark data.\n"
        << "      Currently supports Parquet source files.\n"
        << "      Also supports: --paimon_source_data_file <path>\n"
        << "  --paimon_source_table_path=<path>\n"
        << "      Optional for BM_Read and BM_MOR_Read. If set, read directly from existing\n"
        << "      table path and skip source file loading and pre-write stage.\n"
        << "      Also supports: --paimon_source_table_path <path>\n"
        << "  --paimon_pk_columns=<col1,col2,...>\n"
        << "      Required by BM_PK_Write and BM_MOR_Read.\n"
        << "      Also supports: --paimon_pk_columns <col1,col2,...>\n"
        << "  --paimon_option=<key1>:<value1>;<key2>:<value2>\n"
        << "      Optional and repeatable. Pass through table options as-is.\n"
        << "      Default table file format is parquet; use file.format:<format> to override.\n"
        << "      Also supports: --paimon_option <key1>:<value1>;<key2>:<value2>\n"
        << "      Note: use quotes in shell, e.g. \"--paimon_option k1:v1;k2:v2\".\n"
        << "\n"
        << "Example:\n"
        << "  paimon-read-write-benchmark --paimon_source_data_file /path/data.parquet \\\n"
        << "      --paimon_pk_columns=id --paimon_option \"read.batch-size:8192\" \\\n"
        << "      --benchmark_filter=BM_Read\n"
        << std::endl;
}

Result<std::unique_ptr<paimon::test::UniqueTestDirectory>> CreateBenchmarkWorkspace() {
    auto workspace = paimon::test::UniqueTestDirectory::Create();
    if (workspace == nullptr) {
        return Status::Invalid("failed to create benchmark workspace");
    }
    return workspace;
}

uint64_t NextTableId() {
    static std::atomic<uint64_t> id{0};
    return ++id;
}

std::string RequirePath(const std::string& root_path, const std::string& db_name,
                        const std::string& table_name) {
    return root_path + "/" + db_name + ".db/" + table_name;
}

template <typename T>
Result<T> AddContext(paimon::Result<T>&& result, const std::string& context) {
    if (!result.ok()) {
        const Status status = result.status();
        return status.WithMessage(context, ": ", status.message());
    }
    return std::move(result).value();
}

Status AddContext(const paimon::Status& status, const std::string& context) {
    if (!status.ok()) {
        return status.WithMessage(context, ": ", status.message());
    }
    return Status::OK();
}

void SkipWithMessage(::benchmark::State& state, const std::string& message) {
    state.SkipWithError(message);
}

std::string GetConfiguredFileFormat() {
    std::string file_format = "parquet";
    for (const auto& kv : GetBenchmarkCliOptions().extra_options) {
        if (kv.first == paimon::Options::FILE_FORMAT) {
            file_format = kv.second;
        }
    }
    return file_format;
}

bool IsFileFormatSupported(const std::string& format) {
    if (format == "parquet") {
        return true;
    }
    if (format == "orc") {
#ifdef PAIMON_ENABLE_ORC
        return true;
#else
        return false;
#endif
    }
    return false;
}

void ApplyExtraOptions(std::map<std::string, std::string>* options) {
    for (const auto& kv : GetBenchmarkCliOptions().extra_options) {
        (*options)[kv.first] = kv.second;
    }
}

std::map<std::string, std::string> BuildOptions(const std::string& file_format) {
    std::map<std::string, std::string> options = {
        {paimon::Options::FILE_FORMAT, file_format},
    };
    ApplyExtraOptions(&options);
    return options;
}

std::map<std::string, std::string> BuildPkOptions(const std::string& file_format) {
    auto options = BuildOptions(file_format);
    options[paimon::Options::BUCKET] = "1";
    options[paimon::Options::MERGE_ENGINE] = "deduplicate";
    return options;
}

std::string GetSourceDataFilePath() {
    return GetBenchmarkCliOptions().source_data_file;
}

std::string GetSourceTablePath() {
    return GetBenchmarkCliOptions().source_table_path;
}

const std::vector<std::string>& GetPkColumns() {
    return GetBenchmarkCliOptions().pk_columns;
}

SourceDataSpec GetSourceDataSpec() {
    const std::string source_data_file_path = GetSourceDataFilePath();
    if (!source_data_file_path.empty()) {
        return {"parquet", source_data_file_path};
    }
    return {"", ""};
}

int64_t GetSourceBatchMaxRows() {
    return kSourceBatchMaxRows;
}

int32_t GetRowToBatchThreadNumber() {
    return kRowToBatchThreadNumber;
}

bool SupportsParquetSourceDataMode() {
#if PAIMON_BENCHMARK_HAS_PARQUET_READER
    return true;
#else
    return false;
#endif
}

bool SupportsSourceDataMode(const std::string& source_format) {
    if (source_format == "parquet") {
        return SupportsParquetSourceDataMode();
    }
    return false;
}

struct SourceDataMetadata {
    std::shared_ptr<arrow::Schema> schema;
    int64_t total_rows = 0;
    std::string format;
    std::string path;
};

#if PAIMON_BENCHMARK_HAS_PARQUET_READER
Result<std::unique_ptr<parquet::arrow::FileReader>> OpenParquetSourceReader(
    const std::string& path) {
    auto input = arrow::io::ReadableFile::Open(path);
    if (!input.ok()) {
        return Status::Invalid("open Parquet source failed: ", path, ", ",
                               input.status().ToString());
    }

    std::unique_ptr<parquet::arrow::FileReader> parquet_reader;
    const auto open_status = parquet::arrow::OpenFile(
        input.ValueUnsafe(), arrow::default_memory_pool(), &parquet_reader);
    if (!open_status.ok()) {
        return Status::Invalid("create Parquet reader failed: ", open_status.ToString());
    }
    parquet_reader->set_batch_size(GetSourceBatchMaxRows());
    return parquet_reader;
}
#endif

Result<SourceDataMetadata> LoadParquetSourceMetadata(const std::string& path) {
#if !PAIMON_BENCHMARK_HAS_PARQUET_READER
    return Status::Invalid(
        "Parquet source data mode requires parquet::arrow reader support in this build");
#else
    static SourceDataMetadata cache;
    static std::mutex cache_mutex;
    std::lock_guard<std::mutex> lock(cache_mutex);
    if (cache.path == path && cache.format == "parquet") {
        return cache;
    }

    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<parquet::arrow::FileReader> parquet_reader,
                           OpenParquetSourceReader(path));
    std::shared_ptr<arrow::Schema> schema;
    const auto schema_status = parquet_reader->GetSchema(&schema);
    if (!schema_status.ok()) {
        return Status::Invalid("read Parquet source schema failed: ", schema_status.ToString());
    }

    const int64_t total_rows = parquet_reader->parquet_reader()->metadata()->num_rows();
    if (total_rows <= 0) {
        return Status::Invalid("Parquet source is empty: ", path);
    }

    cache.schema = std::move(schema);
    cache.total_rows = total_rows;
    cache.format = "parquet";
    cache.path = path;
    return cache;
#endif
}

Result<SourceDataMetadata> LoadSourceDataMetadata(const SourceDataSpec& source_spec) {
    if (source_spec.format == "parquet") {
        return LoadParquetSourceMetadata(source_spec.path);
    }
    return Status::Invalid("unknown source format: ", source_spec.format);
}

std::shared_ptr<arrow::StructArray> BuildStructArrayFromRecordBatch(
    const std::shared_ptr<arrow::RecordBatch>& batch) {
    return std::make_shared<arrow::StructArray>(arrow::struct_(batch->schema()->fields()),
                                                batch->num_rows(), batch->columns());
}

Result<std::unique_ptr<paimon::RecordBatch>> MakeRecordBatch(
    const std::shared_ptr<arrow::StructArray>& arr) {
    ArrowArray c_array;
    if (!arrow::ExportArray(*arr, &c_array).ok()) {
        return Status::Invalid("failed to export arrow array");
    }
    paimon::RecordBatchBuilder builder(&c_array);
    return AddContext(builder.Finish(), "build paimon record batch");
}

Status EnsureTable(const std::string& root_path, const std::string& db_name,
                   const std::string& table_name, const std::map<std::string, std::string>& options,
                   const std::shared_ptr<arrow::Schema>& schema,
                   const std::vector<std::string>& primary_keys = {}) {
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::Catalog> catalog,
        AddContext(paimon::Catalog::Create(root_path, options), "create catalog"));
    PAIMON_RETURN_NOT_OK(
        AddContext(catalog->CreateDatabase(db_name, options, true), "create database"));

    ArrowSchema c_schema;
    if (!arrow::ExportSchema(*schema, &c_schema).ok()) {
        return Status::Invalid("failed to export table schema");
    }
    PAIMON_RETURN_NOT_OK(
        AddContext(catalog->CreateTable(paimon::Identifier(db_name, table_name), &c_schema,
                                        /*partition_keys=*/{}, primary_keys, options,
                                        /*ignore_if_exists=*/false),
                   "create table"));
    return Status::OK();
}

Status WriteSourceDataToWriter(paimon::FileStoreWrite* writer, const SourceDataSpec& source_spec) {
    if (source_spec.format != "parquet") {
        return Status::Invalid("unknown source format: ", source_spec.format);
    }

#if !PAIMON_BENCHMARK_HAS_PARQUET_READER
    return Status::Invalid(
        "Parquet source data mode requires parquet::arrow reader support in this build");
#else
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<parquet::arrow::FileReader> parquet_reader,
                           OpenParquetSourceReader(source_spec.path));
    std::unique_ptr<arrow::RecordBatchReader> batch_reader;
    const auto reader_status = parquet_reader->GetRecordBatchReader(&batch_reader);
    if (!reader_status.ok()) {
        return Status::Invalid("create Parquet source batch reader failed: ",
                               reader_status.ToString());
    }

    int64_t written_rows = 0;
    while (true) {
        std::shared_ptr<arrow::RecordBatch> record_batch;
        const auto read_status = batch_reader->ReadNext(&record_batch);
        if (!read_status.ok()) {
            return Status::Invalid("read Parquet source batch failed: ", read_status.ToString());
        }
        if (record_batch == nullptr) {
            break;
        }
        if (record_batch->num_rows() <= 0) {
            continue;
        }

        auto struct_array = BuildStructArrayFromRecordBatch(record_batch);
        PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::RecordBatch> batch,
                               MakeRecordBatch(struct_array));
        PAIMON_RETURN_NOT_OK(AddContext(writer->Write(std::move(batch)), "write batch"));
        written_rows += record_batch->num_rows();
    }

    if (written_rows <= 0) {
        return Status::Invalid("source file has no non-empty data batches: ", source_spec.path);
    }
    return Status::OK();
#endif
}

Status WriteAndCommit(const std::string& table_path,
                      const std::map<std::string, std::string>& options,
                      const SourceDataSpec& source_spec) {
    paimon::WriteContextBuilder write_builder(table_path, "benchmark-writer");
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::WriteContext> write_ctx,
        AddContext(write_builder.SetOptions(options).Finish(), "create write context"));
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::FileStoreWrite> writer,
                           AddContext(paimon::FileStoreWrite::Create(std::move(write_ctx)),
                                      "create file store writer"));

    PAIMON_RETURN_NOT_OK(WriteSourceDataToWriter(writer.get(), source_spec));
    PAIMON_ASSIGN_OR_RAISE(std::vector<std::shared_ptr<paimon::CommitMessage>> messages,
                           AddContext(writer->PrepareCommit(), "prepare commit"));

    paimon::CommitContextBuilder commit_builder(table_path, "benchmark-writer");
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::CommitContext> commit_ctx,
        AddContext(commit_builder.SetOptions(options).Finish(), "create commit context"));
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::FileStoreCommit> committer,
        AddContext(paimon::FileStoreCommit::Create(std::move(commit_ctx)), "create committer"));
    PAIMON_RETURN_NOT_OK(AddContext(committer->Commit(messages), "commit write"));
    return Status::OK();
}

struct SharedReadTableCache {
    std::string key;
    std::unique_ptr<paimon::test::UniqueTestDirectory> workspace;
    std::string table_path;
    int64_t total_rows = 0;
};

struct SharedMorReadTableCache {
    std::string key;
    std::unique_ptr<paimon::test::UniqueTestDirectory> workspace;
    std::string table_path;
    int64_t total_rows = 0;
};

std::string BuildReadTableCacheKey(const std::string& file_format,
                                   const SourceDataSpec& source_spec) {
    return file_format + "|" + source_spec.format + "|" + source_spec.path + "|" +
           std::to_string(GetSourceBatchMaxRows());
}

std::string JoinColumns(const std::vector<std::string>& columns) {
    std::string joined;
    for (size_t i = 0; i < columns.size(); ++i) {
        if (i > 0) {
            joined.append(",");
        }
        joined.append(columns[i]);
    }
    return joined;
}

Result<const SharedMorReadTableCache*> GetOrCreateSharedMorReadTable(
    const std::string& file_format, const SourceDataSpec& source_spec) {
    static SharedMorReadTableCache cache;
    static std::mutex cache_mutex;

    const std::vector<std::string>& pk_columns = GetPkColumns();
    const std::string cache_key =
        BuildReadTableCacheKey(file_format, source_spec) + "|pk=" + JoinColumns(pk_columns);
    std::lock_guard<std::mutex> lock(cache_mutex);
    if (cache.workspace != nullptr && cache.key == cache_key) {
        std::cout << "[benchmark][mor-read] reuse_output_table_path=" << cache.table_path
                  << std::endl;
        return &cache;
    }

    auto options = BuildPkOptions(file_format);
    PAIMON_ASSIGN_OR_RAISE(const SourceDataMetadata source_metadata,
                           LoadSourceDataMetadata(source_spec));

    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::test::UniqueTestDirectory> workspace,
                           CreateBenchmarkWorkspace());
    const std::string db_name = "bench_db";
    const std::string table_name = "mor_read_shared_" + std::to_string(NextTableId());
    PAIMON_RETURN_NOT_OK(EnsureTable(workspace->Str(), db_name, table_name, options,
                                     source_metadata.schema,
                                     /*primary_keys=*/pk_columns));
    const std::string table_path = RequirePath(workspace->Str(), db_name, table_name);
    std::cout << "[benchmark][mor-read] create_shared_output_table_path=" << table_path
              << std::endl;
    PAIMON_RETURN_NOT_OK(WriteAndCommit(table_path, options, source_spec));

    cache.key = cache_key;
    cache.workspace = std::move(workspace);
    cache.table_path = table_path;
    cache.total_rows = source_metadata.total_rows;
    return &cache;
}

Result<const SharedReadTableCache*> GetOrCreateSharedReadTable(const std::string& file_format,
                                                               const SourceDataSpec& source_spec) {
    static SharedReadTableCache cache;
    static std::mutex cache_mutex;

    const std::string cache_key = BuildReadTableCacheKey(file_format, source_spec);
    std::lock_guard<std::mutex> lock(cache_mutex);
    if (cache.workspace != nullptr && cache.key == cache_key) {
        std::cout << "[benchmark][read] reuse_output_table_path=" << cache.table_path << std::endl;
        return &cache;
    }

    auto options = BuildOptions(file_format);
    PAIMON_ASSIGN_OR_RAISE(const SourceDataMetadata source_metadata,
                           LoadSourceDataMetadata(source_spec));

    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::test::UniqueTestDirectory> workspace,
                           CreateBenchmarkWorkspace());
    const std::string db_name = "bench_db";
    const std::string table_name = "read_shared_" + std::to_string(NextTableId());
    PAIMON_RETURN_NOT_OK(
        EnsureTable(workspace->Str(), db_name, table_name, options, source_metadata.schema));
    const std::string table_path = RequirePath(workspace->Str(), db_name, table_name);
    std::cout << "[benchmark][read] create_shared_output_table_path=" << table_path << std::endl;
    PAIMON_RETURN_NOT_OK(WriteAndCommit(table_path, options, source_spec));

    cache.key = cache_key;
    cache.workspace = std::move(workspace);
    cache.table_path = table_path;
    cache.total_rows = source_metadata.total_rows;
    return &cache;
}

Result<int64_t> ReadRows(const std::string& table_path,
                         const std::map<std::string, std::string>& options,
                         int32_t prefetch_parallel_num) {
    paimon::ScanContextBuilder scan_builder(table_path);
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::ScanContext> scan_ctx,
        AddContext(scan_builder.SetOptions(options).Finish(), "create scan context"));
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::TableScan> scanner,
        AddContext(paimon::TableScan::Create(std::move(scan_ctx)), "create scanner"));
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<paimon::Plan> plan,
                           AddContext(scanner->CreatePlan(), "create plan"));

    paimon::ReadContextBuilder read_builder(table_path);
    constexpr int32_t kPrefetchBatchCount = 600;
    read_builder.SetOptions(options)
        .EnablePrefetch(true)
        .SetPrefetchBatchCount(kPrefetchBatchCount)
        .SetPrefetchMaxParallelNum(prefetch_parallel_num)
        .EnableMultiThreadRowToBatch(GetRowToBatchThreadNumber() > 1)
        .SetRowToBatchThreadNumber(GetRowToBatchThreadNumber());
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::ReadContext> read_ctx,
                           AddContext(read_builder.Finish(), "create read context"));
    PAIMON_ASSIGN_OR_RAISE(
        std::unique_ptr<paimon::TableRead> reader,
        AddContext(paimon::TableRead::Create(std::move(read_ctx)), "create table reader"));
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<paimon::BatchReader> batch_reader,
                           AddContext(reader->CreateReader(plan->Splits()), "create batch reader"));

    int64_t total_rows = 0;
    while (true) {
        PAIMON_ASSIGN_OR_RAISE(paimon::BatchReader::ReadBatch batch,
                               AddContext(batch_reader->NextBatch(), "read next batch"));
        if (paimon::BatchReader::IsEofBatch(batch)) {
            break;
        }
        auto& [array, schema] = batch;
        auto imported = arrow::ImportArray(array.get(), schema.get());
        if (!imported.ok()) {
            return Status::Invalid("import c data array failed: ", imported.status().ToString());
        }
        total_rows += imported.ValueUnsafe()->length();
    }

    return total_rows;
}

struct PreparedSourceData {
    std::shared_ptr<arrow::Schema> schema;
    int64_t total_rows = 0;
};

bool TryGetSourceSpec(::benchmark::State& state, SourceDataSpec* source_spec) {
    (void)state;
    *source_spec = GetSourceDataSpec();
    return true;
}

bool TryPrepareSourceData(::benchmark::State& state, const SourceDataSpec& source_spec,
                          PreparedSourceData* prepared) {
    auto source_metadata = LoadSourceDataMetadata(source_spec);
    if (!source_metadata.ok()) {
        SkipWithMessage(state, source_metadata.status().ToString());
        return false;
    }
    prepared->schema = source_metadata.value().schema;
    prepared->total_rows = source_metadata.value().total_rows;
    return true;
}

}  // namespace

Status ParsePaimonBenchmarkCliArgs(int* argc, char** argv) {
    auto parsed_argc = static_cast<int32_t>(*argc);
    PAIMON_RETURN_NOT_OK(ParsePaimonBenchmarkCliArgsImpl(&parsed_argc, argv));
    *argc = static_cast<int>(parsed_argc);
    return Status::OK();
}

bool HasHelpFlag(int32_t argc, char** argv) {
    return HasHelpFlagImpl(argc, argv);
}

void PrintPaimonBenchmarkCliHelp() {
    PrintPaimonBenchmarkCliHelpImpl();
}

void RunBMWrite(::benchmark::State& state) {
    const std::string file_format = GetConfiguredFileFormat();
    SourceDataSpec source_spec;
    if (!TryGetSourceSpec(state, &source_spec)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip(
            state, source_spec.path, "--paimon_source_data_file is required", &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format,
                                                       SupportsSourceDataMode(source_spec.format),
                                                       &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateFileFormatOrSkip(
            state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) {
        return;
    }

    auto options = BuildOptions(file_format);
    PreparedSourceData prepared;
    if (!TryPrepareSourceData(state, source_spec, &prepared)) {
        return;
    }
    auto workspace = CreateBenchmarkWorkspace();
    if (!workspace.ok()) {
        SkipWithMessage(state, workspace.status().ToString());
        return;
    }

    for (auto _ : state) {
        const std::string db_name = "bench_db";
        const std::string table_name = "write_" + std::to_string(NextTableId());
        const Status ensure_status =
            EnsureTable(workspace.value()->Str(), db_name, table_name, options, prepared.schema);
        if (!ensure_status.ok()) {
            SkipWithMessage(state, ensure_status.ToString());
            return;
        }
        const std::string table_path = RequirePath(workspace.value()->Str(), db_name, table_name);
        std::cout << "[benchmark][write] output_table_path=" << table_path << std::endl;
        const Status write_status = WriteAndCommit(table_path, options, source_spec);
        if (!write_status.ok()) {
            SkipWithMessage(state, write_status.ToString());
            return;
        }
    }

    state.SetItemsProcessed(state.iterations() * prepared.total_rows);
}

void RunBMRead(::benchmark::State& state) {
    const auto prefetch_parallel_num = static_cast<int32_t>(state.range(0));
    const std::string file_format = GetConfiguredFileFormat();
    const std::string source_table_path = GetSourceTablePath();
    SourceDataSpec source_spec;
    if (!TryGetSourceSpec(state, &source_spec)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateFileFormatOrSkip(
            state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) {
        return;
    }

    if (!BenchmarkHelpers::ValidatePrefetchParallelOrSkip(state, prefetch_parallel_num,
                                                          &SkipWithMessage)) {
        return;
    }

    auto options = BuildOptions(file_format);

    auto source_table_read_result = BenchmarkHelpers::TryRunSourceTableReadMode(
        state, "read", source_table_path,
        [&]() { return ReadRows(source_table_path, options, prefetch_parallel_num); });
    if (!source_table_read_result.ok()) {
        SkipWithMessage(state, source_table_read_result.status().ToString());
        return;
    }
    if (source_table_read_result.value()) {
        return;
    }

    if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip(
            state, source_spec.path,
            "--paimon_source_data_file is required when --paimon_source_table_path is not set",
            &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format,
                                                       SupportsSourceDataMode(source_spec.format),
                                                       &SkipWithMessage)) {
        return;
    }

    auto shared_table = GetOrCreateSharedReadTable(file_format, source_spec);
    if (!shared_table.ok()) {
        SkipWithMessage(state, shared_table.status().ToString());
        return;
    }

    auto rows_read = BenchmarkHelpers::RunReadIterations(state, [&]() {
        return ReadRows(shared_table.value()->table_path, options, prefetch_parallel_num);
    });
    if (!rows_read.ok()) {
        SkipWithMessage(state, rows_read.status().ToString());
        return;
    }

    state.SetItemsProcessed(state.iterations() * rows_read.value());
}

void RunBMPkWrite(::benchmark::State& state) {
    const std::string file_format = GetConfiguredFileFormat();
    SourceDataSpec source_spec;
    if (!TryGetSourceSpec(state, &source_spec)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip(
            state, source_spec.path, "--paimon_source_data_file is required", &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format,
                                                       SupportsSourceDataMode(source_spec.format),
                                                       &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateFileFormatOrSkip(
            state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) {
        return;
    }
    const std::vector<std::string>& pk_columns = GetPkColumns();
    if (pk_columns.empty()) {
        SkipWithMessage(state, "--paimon_pk_columns is required for BM_PK_Write");
        return;
    }

    auto options = BuildPkOptions(file_format);
    PreparedSourceData prepared;
    if (!TryPrepareSourceData(state, source_spec, &prepared)) {
        return;
    }
    auto workspace = CreateBenchmarkWorkspace();
    if (!workspace.ok()) {
        SkipWithMessage(state, workspace.status().ToString());
        return;
    }

    for (auto _ : state) {
        const std::string db_name = "bench_db";
        const std::string table_name = "pk_write_" + std::to_string(NextTableId());
        const Status ensure_status =
            EnsureTable(workspace.value()->Str(), db_name, table_name, options, prepared.schema,
                        /*primary_keys=*/pk_columns);
        if (!ensure_status.ok()) {
            SkipWithMessage(state, ensure_status.ToString());
            return;
        }
        const std::string table_path = RequirePath(workspace.value()->Str(), db_name, table_name);
        std::cout << "[benchmark][pk-write] output_table_path=" << table_path << std::endl;
        const Status write_status = WriteAndCommit(table_path, options, source_spec);
        if (!write_status.ok()) {
            SkipWithMessage(state, write_status.ToString());
            return;
        }
    }

    state.SetItemsProcessed(state.iterations() * prepared.total_rows);
}

void RunBMMorRead(::benchmark::State& state) {
    const auto prefetch_parallel_num = static_cast<int32_t>(state.range(0));
    const std::string file_format = GetConfiguredFileFormat();
    const std::string source_table_path = GetSourceTablePath();
    SourceDataSpec source_spec;
    if (!TryGetSourceSpec(state, &source_spec)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateFileFormatOrSkip(
            state, file_format, IsFileFormatSupported(file_format), &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidatePrefetchParallelOrSkip(state, prefetch_parallel_num,
                                                          &SkipWithMessage)) {
        return;
    }

    const auto source_table_read_options = BuildOptions(file_format);
    auto source_table_read_result =
        BenchmarkHelpers::TryRunSourceTableReadMode(state, "mor-read", source_table_path, [&]() {
            return ReadRows(source_table_path, source_table_read_options, prefetch_parallel_num);
        });
    if (!source_table_read_result.ok()) {
        SkipWithMessage(state, source_table_read_result.status().ToString());
        return;
    }
    if (source_table_read_result.value()) {
        return;
    }

    if (!BenchmarkHelpers::ValidateSourcePresenceOrSkip(
            state, source_spec.path,
            "--paimon_source_data_file is required when --paimon_source_table_path is not set",
            &SkipWithMessage)) {
        return;
    }
    if (!BenchmarkHelpers::ValidateSourceSupportOrSkip(state, source_spec.format,
                                                       SupportsSourceDataMode(source_spec.format),
                                                       &SkipWithMessage)) {
        return;
    }
    if (GetPkColumns().empty()) {
        SkipWithMessage(state, "--paimon_pk_columns is required for BM_MOR_Read");
        return;
    }

    auto options = BuildPkOptions(file_format);
    auto shared_table = GetOrCreateSharedMorReadTable(file_format, source_spec);
    if (!shared_table.ok()) {
        SkipWithMessage(state, shared_table.status().ToString());
        return;
    }

    auto rows_read = BenchmarkHelpers::RunReadIterations(state, [&]() {
        return ReadRows(shared_table.value()->table_path, options, prefetch_parallel_num);
    });
    if (!rows_read.ok()) {
        SkipWithMessage(state, rows_read.status().ToString());
        return;
    }
    state.SetItemsProcessed(state.iterations() * rows_read.value());
}

}  // namespace paimon::benchmark