/*
 * 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/core/operation/append_only_file_store_write.h"

#include <functional>
#include <limits>
#include <vector>

#include "arrow/c/bridge.h"
#include "arrow/c/helpers.h"
#include "paimon/common/data/binary_row.h"
#include "paimon/common/data/shredding/shredding_write_plan_factories.h"
#include "paimon/common/table/special_fields.h"
#include "paimon/common/utils/arrow/arrow_utils.h"
#include "paimon/core/append/append_only_writer.h"
#include "paimon/core/append/bucketed_append_compact_manager.h"
#include "paimon/core/compact/noop_compact_manager.h"
#include "paimon/core/core_options.h"
#include "paimon/core/io/append_data_file_writer_factory.h"
#include "paimon/core/io/data_file_meta.h"
#include "paimon/core/io/data_file_path_factory.h"
#include "paimon/core/io/rolling_file_writer.h"
#include "paimon/core/io/shredding_append_data_file_writer_factory.h"
#include "paimon/core/manifest/manifest_file.h"
#include "paimon/core/manifest/manifest_list.h"
#include "paimon/core/operation/append_only_file_store_scan.h"
#include "paimon/core/operation/file_store_scan.h"
#include "paimon/core/operation/internal_read_context.h"
#include "paimon/core/operation/raw_file_split_read.h"
#include "paimon/core/operation/restore_files.h"
#include "paimon/core/schema/table_schema.h"
#include "paimon/core/snapshot.h"
#include "paimon/core/utils/file_store_path_factory.h"
#include "paimon/core/utils/snapshot_manager.h"
#include "paimon/executor.h"
#include "paimon/logging.h"
#include "paimon/read_context.h"
#include "paimon/result.h"
namespace arrow {
class Schema;
}  // namespace arrow

namespace paimon {
class DataFilePathFactory;
class MemoryPool;
class SchemaManager;

AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite(
    const std::shared_ptr<FileStorePathFactory>& file_store_path_factory,
    const std::shared_ptr<SnapshotManager>& snapshot_manager,
    const std::shared_ptr<SchemaManager>& schema_manager, const std::string& commit_user,
    const std::string& root_path, const std::shared_ptr<TableSchema>& table_schema,
    const std::shared_ptr<arrow::Schema>& schema,
    const std::shared_ptr<arrow::Schema>& write_schema,
    const std::shared_ptr<arrow::Schema>& partition_schema,
    const std::shared_ptr<BucketedDvMaintainer::Factory>& dv_maintainer_factory,
    const std::shared_ptr<IOManager>& io_manager, const CoreOptions& options,
    bool ignore_previous_files, bool is_streaming_mode, bool ignore_num_bucket_check,
    const std::shared_ptr<Executor>& executor, const std::shared_ptr<MemoryPool>& pool)
    : AbstractFileStoreWrite(file_store_path_factory, snapshot_manager, schema_manager, commit_user,
                             root_path, table_schema, schema, write_schema, partition_schema,
                             dv_maintainer_factory, io_manager, options, ignore_previous_files,
                             is_streaming_mode, ignore_num_bucket_check, executor, pool),
      logger_(Logger::GetLogger("AppendOnlyFileStoreWrite")) {
    write_cols_ = write_schema->field_names();
    // optimize write_cols to null in following cases:
    // 1. write_schema contains all columns
    // 2. TODO(xinyu.lxy) write_schema contains all columns and append _ROW_ID & _SEQUENCE_NUMBER
    // cols
    if (schema->Equals(write_schema)) {
        write_cols_ = std::nullopt;
    }
}

AppendOnlyFileStoreWrite::~AppendOnlyFileStoreWrite() = default;

Result<std::unique_ptr<FileStoreScan>> AppendOnlyFileStoreWrite::CreateFileStoreScan(
    const std::shared_ptr<ScanFilter>& scan_filter) const {
    PAIMON_ASSIGN_OR_RAISE(
        std::shared_ptr<ManifestList> manifest_list,
        ManifestList::Create(options_.GetFileSystem(), options_.GetManifestFormat(),
                             options_.GetManifestCompression(), file_store_path_factory_,
                             options_.GetCache(), pool_));
    PAIMON_ASSIGN_OR_RAISE(
        std::shared_ptr<ManifestFile> manifest_file,
        ManifestFile::Create(options_.GetFileSystem(), options_.GetManifestFormat(),
                             options_.GetManifestCompression(), file_store_path_factory_,
                             options_.GetManifestTargetFileSize(), pool_, options_,
                             partition_schema_));
    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileStoreScan> scan,
                           AppendOnlyFileStoreScan::Create(
                               snapshot_manager_, schema_manager_, manifest_list, manifest_file,
                               table_schema_, schema_, scan_filter, options_, executor_, pool_));
    return scan;
}

Result<std::vector<std::shared_ptr<DataFileMeta>>> AppendOnlyFileStoreWrite::CompactRewrite(
    const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory,
    const std::vector<std::shared_ptr<DataFileMeta>>& to_compact,
    const std::shared_ptr<CancellationController>& cancellation_controller) {
    if (to_compact.empty()) {
        return std::vector<std::shared_ptr<DataFileMeta>>{};
    }

    PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<BatchReader> reader,
                           CreateFilesReader(partition, bucket, dv_factory, to_compact));
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataFilePathFactory> data_file_path_factory,
                           file_store_path_factory_->CreateDataFilePathFactory(partition, bucket));
    PAIMON_ASSIGN_OR_RAISE(
        WriterFactory writer_factory,
        GetDataFileWriterFactory(data_file_path_factory, write_schema_, write_cols_, to_compact));
    auto rewriter =
        std::make_unique<RollingFileWriter<::ArrowArray*, std::shared_ptr<DataFileMeta>>>(
            options_.GetTargetFileSize(/*has_primary_key=*/false),
            /*target_file_row_num=*/std::numeric_limits<int64_t>::max(), writer_factory);

    ScopeGuard reader_guard([&]() {
        if (reader) {
            reader->Close();
        }
    });

    ScopeGuard rewriter_guard([&]() {
        if (rewriter) {
            (void)rewriter->Close();
        }
    });

    while (true) {
        if (cancellation_controller->IsCancelled()) {
            return Status::Cancelled("Compaction cancelled while rewriting files.");
        }
        PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, reader->NextBatch());
        if (BatchReader::IsEofBatch(batch)) {
            break;
        }
        auto& [c_array, c_schema] = batch;
        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> arrow_array,
                                          arrow::ImportArray(c_array.get(), c_schema.get()));
        auto struct_array = std::dynamic_pointer_cast<arrow::StructArray>(arrow_array);
        if (!struct_array) {
            return Status::Invalid(
                "cannot cast array to StructArray in CompleteRowKindBatchReader");
        }
        PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray(
                                                 struct_array, SpecialFields::ValueKind().Name()));
        PAIMON_RETURN_NOT_OK_FROM_ARROW(
            arrow::ExportArray(*struct_array, c_array.get(), c_schema.get()));
        ArrowSchemaRelease(c_schema.get());
        ScopeGuard guard([array = c_array.get()]() { ArrowArrayRelease(array); });
        PAIMON_RETURN_NOT_OK(rewriter->Write(c_array.get()));
        guard.Release();
    }
    rewriter_guard.Release();
    PAIMON_RETURN_NOT_OK(rewriter->Close());
    return rewriter->GetResult();
}

Result<std::shared_ptr<BatchWriter>> AppendOnlyFileStoreWrite::CreateWriter(
    const BinaryRow& partition, int32_t bucket,
    const std::vector<std::shared_ptr<DataFileMeta>>& restore_data_files,
    int64_t restore_max_seq_number, const std::shared_ptr<BucketedDvMaintainer>& dv_maintainer) {
    PAIMON_LOG_DEBUG(logger_, "Creating append only writer for partition %s, bucket %d",
                     partition.ToString().c_str(), bucket);
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<DataFilePathFactory> data_file_path_factory,
                           file_store_path_factory_->CreateDataFilePathFactory(partition, bucket));

    std::shared_ptr<CompactManager> compact_manager;
    if (options_.WriteOnly() || options_.DataEvolutionEnabled() || options_.GetBucket() == -1) {
        compact_manager = std::make_shared<NoopCompactManager>();
    } else {
        auto dv_factory =
            [dv_maintainer](
                const std::string& file_name) -> Result<std::shared_ptr<DeletionVector>> {
            if (dv_maintainer) {
                return dv_maintainer->DeletionVectorOf(file_name).value_or(
                    std::shared_ptr<DeletionVector>());
            }
            return std::shared_ptr<DeletionVector>();
        };
        auto cancellation_controller = std::make_shared<CancellationController>();

        auto rewriter = [this, partition, bucket, dv_factory, cancellation_controller](
                            const std::vector<std::shared_ptr<DataFileMeta>>& to_compact)
            -> Result<std::vector<std::shared_ptr<DataFileMeta>>> {
            return CompactRewrite(partition, bucket, dv_factory, to_compact,
                                  cancellation_controller);
        };

        compact_manager = std::make_shared<BucketedAppendCompactManager>(
            compact_executor_, restore_data_files, dv_maintainer,
            options_.GetCompactionMinFileNum(),
            options_.GetTargetFileSize(/*has_primary_key=*/false),
            options_.GetCompactionFileSize(/*has_primary_key=*/false),
            options_.CompactionForceRewriteAllFiles(), rewriter,
            compaction_metrics_->CreateReporter(partition, bucket), cancellation_controller);
    }

    auto writer = std::make_unique<AppendOnlyWriter>(
        options_, table_schema_->Id(), write_schema_, write_cols_, restore_max_seq_number,
        data_file_path_factory, compact_manager, pool_);
    return std::shared_ptr<BatchWriter>(std::move(writer));
}

Result<AppendOnlyFileStoreWrite::WriterFactory> AppendOnlyFileStoreWrite::GetDataFileWriterFactory(
    const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
    const std::shared_ptr<arrow::Schema>& schema,
    const std::optional<std::vector<std::string>>& write_cols,
    const std::vector<std::shared_ptr<DataFileMeta>>& to_compact) const {
    auto seq_num_counter = std::make_shared<LongCounter>(to_compact[0]->min_sequence_number);
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<ShreddingWritePlanFactory> plan_factory,
                           ShreddingWritePlanFactories::SelectActive(options_, schema, pool_));
    if (plan_factory != nullptr) {
        return std::make_shared<ShreddingAppendDataFileWriterFactory>(
            options_, table_schema_->Id(), schema, write_cols, seq_num_counter,
            FileSource::Compact(), data_file_path_factory, plan_factory, pool_);
    }
    return std::make_shared<AppendDataFileWriterFactory>(
        options_, table_schema_->Id(), schema, write_cols, seq_num_counter, FileSource::Compact(),
        data_file_path_factory, pool_);
}

Result<std::unique_ptr<BatchReader>> AppendOnlyFileStoreWrite::CreateFilesReader(
    const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory,
    const std::vector<std::shared_ptr<DataFileMeta>>& files) const {
    ReadContextBuilder context_builder(root_path_);
    context_builder.SetOptions(options_.ToMap())
        .WithFileSystem(options_.GetFileSystem())
        .EnablePrefetch(true)
        .SetPrefetchMaxParallelNum(1)
        .SetPrefetchBatchCount(3)
        .WithMemoryPool(pool_);
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<ReadContext> read_context, context_builder.Finish());
    std::map<std::string, std::string> options = options_.ToMap();
    // TODO(xinyu.lxy): temporarily disabled pre-buffer for parquet, which may cause high
    // memory usage during compaction. Will fix via parquet format refactor.
    auto new_options = options;
    if (new_options.find("parquet.read.enable-pre-buffer") == new_options.end()) {
        new_options["parquet.read.enable-pre-buffer"] = "false";
    }
    PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<InternalReadContext> internal_read_context,
                           InternalReadContext::Create(read_context, table_schema_, new_options));
    auto read = std::make_unique<RawFileSplitRead>(file_store_path_factory_, internal_read_context,
                                                   pool_, compact_executor_);

    return read->CreateReader(partition, bucket, files, dv_factory);
}

}  // namespace paimon
