// 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.

#pragma once

#include <gen_cpp/DataSinks_types.h>

#include <gen_cpp/PlanNodes_types.h>

#include <map>

#include <memory>

#include <string>

#include <vector>

#include "common/status.h"

#include "core/block/block.h"

#include "exec/sink/writer/async_result_writer.h"

#include "exec/sink/writer/iceberg/viceberg_delete_file_writer.h"

#include "exprs/vexpr_fwd.h"

#include "roaring/roaring64map.hh"

#include "runtime/runtime_profile.h"

namespace doris {

class ObjectPool;

class RuntimeState;

class Dependency;

class VIcebergDeleteFileWriter;

struct IcebergFileDeletion {

    IcebergFileDeletion() = default;

    IcebergFileDeletion(int32_t spec_id, std::string data_json)

            : partition_spec_id(spec_id), partition_data_json(std::move(data_json)) {}

    int32_t partition_spec_id = 0;

    std::string partition_data_json;

    roaring::Roaring64Map rows_to_delete;

};

/**

 * Sink for writing Iceberg position delete files.

 *

 * This sink receives blocks containing a $row_id column with

 * (file_path, row_position, partition_spec_id, partition_data).

 * It groups delete records by file_path and writes position delete files.

 */

class VIcebergDeleteSink final : public AsyncResultWriter {

public:

    VIcebergDeleteSink(const TDataSink& t_sink, const VExprContextSPtrs& output_exprs,

                       std::shared_ptr<Dependency> dep, std::shared_ptr<Dependency> fin_dep);

    ~VIcebergDeleteSink() override = default;

    Status init_properties(ObjectPool* pool);

    Status open(RuntimeState* state, RuntimeProfile* profile) override;

    Status write(RuntimeState* state, Block& block) override;

    Status close(Status) override;

private:

    /**

     * Extract $row_id column from block and group by file_path.

     */

    Status _collect_position_deletes(const Block& block,

                                     std::map<std::string, IcebergFileDeletion>& file_deletions);

    /**

     * Write grouped position deletes to delete files

     */

    Status _write_position_delete_files(

            const std::map<std::string, IcebergFileDeletion>& file_deletions);

    Status _write_deletion_vector_files(

            const std::map<std::string, IcebergFileDeletion>& file_deletions);

    struct DeletionVectorBlob {

        std::string referenced_data_file;

        int32_t partition_spec_id = 0;

        std::string partition_data_json;

        int64_t delete_count = 0; // The number of rows deleted in this delete operation.

        int64_t merged_count =

                0; // The number of rows after merging the old deletion vector and position delete.

        int64_t content_offset = 0;

        int64_t content_size_in_bytes = 0;

        std::vector<char> blob_data;

    };

    Status _write_puffin_file(const std::string& puffin_path,

                              std::vector<DeletionVectorBlob>* blobs, int64_t* out_file_size);

    std::string _build_puffin_footer_json(const std::vector<DeletionVectorBlob>& blobs);

    std::string _generate_puffin_file_path();

    /**

     * Generate unique delete file path

     */

    std::string _generate_delete_file_path(const std::string& referenced_data_file = "");

    /**

     * Get $row_id column index from block

     */

    int _get_row_id_column_index(const Block& block);

    /**

     * Build a block for position delete (file_path, pos)

     */

    Status _build_position_delete_block(const std::string& file_path,

                                        const std::vector<int64_t>& positions, Block& output_block);

    Status _init_position_delete_output_exprs();

    std::string _get_file_extension() const;

    TDataSink _t_sink;

    RuntimeState* _state = nullptr;

    int32_t _format_version = 2;

    TFileContent::type _delete_type = TFileContent::POSITION_DELETES;

    // Writers for delete files

    std::vector<std::unique_ptr<VIcebergDeleteFileWriter>> _writers;

    // Collected commit data from all writers

    std::vector<TIcebergCommitData> _commit_data_list;

    // TODO: All deletions are held in memory until close(). Consider flushing

    //  per-file when the upstream guarantees file_path ordering, or flushing

    //  when estimated memory exceeds a threshold, to reduce peak memory usage.

    std::map<std::string, IcebergFileDeletion> _file_deletions;

    std::map<std::string, std::vector<TIcebergDeleteFileDesc>> _rewritable_delete_files;

    // Hadoop configuration

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

    // File format settings

    TFileFormatType::type _file_format_type = TFileFormatType::FORMAT_PARQUET;

    TFileCompressType::type _compress_type = TFileCompressType::SNAPPYBLOCK;

    // Output directory for delete files

    std::string _output_path;

    std::string _table_location;

    TFileType::type _file_type = TFileType::FILE_HDFS;

    std::vector<TNetworkAddress> _broker_addresses;

    // Partition information

    int32_t _partition_spec_id = 0;

    std::string _partition_data_json;

    // Counters

    size_t _row_count = 0;

    size_t _delete_file_count = 0;

    // Profile counters

    RuntimeProfile::Counter* _written_rows_counter = nullptr;

    RuntimeProfile::Counter* _send_data_timer = nullptr;

    RuntimeProfile::Counter* _write_delete_files_timer = nullptr;

    RuntimeProfile::Counter* _delete_file_count_counter = nullptr;

    RuntimeProfile::Counter* _open_timer = nullptr;

    RuntimeProfile::Counter* _close_timer = nullptr;

    VExprContextSPtrs _position_delete_output_expr_ctxs;

};

} // namespace doris