// 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/olap_file.pb.h>

#include <glog/logging.h>

#include <functional>

#include <memory>

#include <mutex>

#include <optional>

#include <string_view>

#include <unordered_map>

#include <vector>

#include "cloud/config.h"

#include "common/status.h"

#include "io/fs/encrypted_fs_factory.h"

#include "io/fs/file_system.h"

#include "io/fs/file_writer.h"

#include "io/fs/local_file_system.h"

#include "io/fs/packed_file_system.h"

#include "runtime/exec_env.h"

#include "storage/binlog.h"

#include "storage/olap_define.h"

#include "storage/partial_update_info.h"

#include "storage/segment/historical_row_retriever.h"

#include "storage/storage_policy.h"

#include "storage/tablet/tablet.h"

#include "storage/tablet/tablet_schema.h"

namespace doris {

class RowsetWriterContextBuilder;

using RowsetWriterContextBuilderSharedPtr = std::shared_ptr<RowsetWriterContextBuilder>;

class DataDir;

class Tablet;

class FileWriterCreator;

class SegmentCollector;

namespace segment_v2 {

struct HistoricalRowRetrieverContext;

}

struct RowsetWriterContext {

    RowsetWriterContext() : schema_lock(new std::mutex) {

        load_id.set_hi(0);

        load_id.set_lo(0);

    }

    RowsetId rowset_id;

    int64_t db_id {0};

    int64_t table_id {0};

    int64_t tablet_id {0};

    int32_t tablet_schema_hash {0};

    int64_t index_id {0};

    int64_t partition_id {0};

    RowsetTypePB rowset_type {BETA_ROWSET};

    TabletSchemaSPtr tablet_schema;

    // PREPARED/COMMITTED for pending rowset

    // VISIBLE for non-pending rowset

    RowsetStatePB rowset_state {PREPARED};

    // properties for non-pending rowset

    Version version {0, 0};

    // properties for pending rowset

    int64_t txn_id {0};

    int64_t txn_expiration {0}; // For cloud mode

    PUniqueId load_id;

    TabletUid tablet_uid {0, 0};

    // indicate whether the data among segments is overlapping.

    // default is OVERLAP_UNKNOWN.

    SegmentsOverlapPB segments_overlap {OVERLAP_UNKNOWN};

    // segment file use uint32 to represent row number, therefore the maximum is UINT32_MAX.

    // the default is set to INT32_MAX to avoid overflow issue when casting from uint32_t to int.

    // test cases can change this value to control flush timing

    uint32_t max_rows_per_segment = INT32_MAX;

    // not owned, point to the data dir of this rowset

    // for checking disk capacity when write data to disk.

    // ATTN: not support for RowsetConvertor.

    // (because it hard to refactor, and RowsetConvertor will be deprecated in future)

    DataDir* data_dir = nullptr;

    int64_t newest_write_timestamp = -1;

    bool enable_unique_key_merge_on_write = false;

    // store column_unique_id to do index compaction

    std::set<int32_t> columns_to_do_index_compaction;

    // SNII only: (column_unique_id, index_id) pairs whose postings are produced

    // by index compaction. The segment writer raw-builds every OTHER SNII index

    // of the column, so one eligible and one new index on the same column can

    // coexist in a single pass. V2/V3 keep columns_to_do_index_compaction:

    // their per-column CLucene directories cannot split an index off a column.

    std::set<std::pair<int32_t, int64_t>> snii_indexes_to_do_compaction;

    DataWriteType write_type = DataWriteType::TYPE_DEFAULT;

    // need to figure out the sub type of compaction

    ReaderType compaction_type = ReaderType::UNKNOWN;

    BaseTabletSPtr tablet = nullptr;

    std::shared_ptr<MowContext> mow_context;

    std::shared_ptr<FileWriterCreator> file_writer_creator;

    std::shared_ptr<SegmentCollector> segment_collector;

    // memtable_on_sink_support_index_v2 = true, we will create SinkFileWriter to send inverted index file

    bool memtable_on_sink_support_index_v2 = false;

    /// begin file cache opts

    bool write_file_cache = false;

    bool is_hot_data = false;

    uint64_t file_cache_ttl_sec = 0;

    uint64_t approximate_bytes_to_write = 0;

    // If true, compaction output only writes index files to file cache, not data files

    bool compaction_output_write_index_only = false;

    /// end file cache opts

    // segcompaction for this RowsetWriter, only enabled when importing data

    bool enable_segcompaction = false;

    std::shared_ptr<PartialUpdateInfo> partial_update_info;

    bool is_transient_rowset_writer = false;

    segment_v2::HistoricalRowRetrieverContext make_historical_row_retriever_context();

    // Intent flag: caller can actively turn merge-file feature on/off for this rowset.

    // This describes whether we *want* to try small-file merging.

    bool allow_packed_file = true;

    // Physical id of the first segment in this rowset. It can be nonzero for writers that

    // allocate segment ids from a configured range.

    int64_t first_segment_id = 0;

    // Effective flag: whether this context actually ends up using MergeFileSystem for writes.

    // This is decided inside fs() based on enable_merge_file plus other conditions

    // (cloud mode, S3 filesystem, V1 inverted index, global config, etc.), and once

    // set to true it remains stable even if config::enable_merge_file changes later.

    mutable bool packed_file_active = false;

    // Cached FileSystem instance to ensure consistency across multiple fs() calls.

    // This prevents creating multiple MergeFileSystem instances and ensures

    // packed_file_active flag remains consistent.

    mutable io::FileSystemSPtr _cached_fs = nullptr;

    void set_first_segment_id(int64_t segment_id) {

        DORIS_CHECK_GE(segment_id, 0);

        DORIS_CHECK(_cached_fs == nullptr);

        first_segment_id = segment_id;

    }

    // For collect segment statistics for compaction

    std::vector<RowsetReaderSharedPtr> input_rs_readers;

    // TODO(lihangyu) remove this lock

    // In semi-structure senario tablet_schema will be updated concurrently,

    // this lock need to be held when update.Use shared_ptr to avoid delete copy contructor

    std::shared_ptr<std::mutex> schema_lock;

    int64_t compaction_level = 0;

    // For local rowset

    std::string tablet_path;

    // For remote rowset

    std::optional<StorageResource> storage_resource;

    std::optional<EncryptionAlgorithmPB> encrypt_algorithm;

    std::string job_id;

    // Per-segment LSNs allocated before memtable flush. The same storage feeds

    // both the base row LSN column and row-binlog LSN column.

    std::shared_ptr<segment_v2::SegmentAllocatedLsnMap> allocated_lsn_map = nullptr;

    bool _need_allocate_lsn = false;

    void insert_segment_allocated_lsns(int64_t segment_id,

                                       ConstAllocatedLsnVectorSharedPtr allocated_lsns) {

        DCHECK(allocated_lsn_map != nullptr);

        allocated_lsn_map->insert_segment_allocated_lsns(segment_id, std::move(allocated_lsns));

    }

    void remove_segment_allocated_lsns(int64_t segment_id) {

        DCHECK(allocated_lsn_map != nullptr);

        allocated_lsn_map->remove_segment(segment_id);

    }

    ConstAllocatedLsnVectorSharedPtr get_segment_allocated_lsns(int64_t segment_id) const {

        DCHECK(allocated_lsn_map != nullptr);

        return allocated_lsn_map->get_segment_allocated_lsns(segment_id);

    }

    bool need_allocated_lsn() const { return _need_allocate_lsn; }

    bool is_local_rowset() const { return !storage_resource; }

    std::string segment_path(int seg_id) const {

        if (is_local_rowset()) {

            return local_segment_path(tablet_path, rowset_id.to_string(), seg_id);

        } else {

            return storage_resource->remote_segment_path(tablet_id, rowset_id.to_string(), seg_id);

        }

    }

    io::FileSystemSPtr fs() const {

        // Return cached instance if available to ensure consistency across multiple calls

        if (_cached_fs != nullptr) {

            return _cached_fs;

        }

        auto fs = [this]() -> io::FileSystemSPtr {

            if (is_local_rowset()) {

                return io::global_local_filesystem();

            } else {

                return storage_resource->fs;

            }

        }();

        bool is_s3_fs = fs->type() == io::FileSystemType::S3;

        auto algorithm = encrypt_algorithm;

        if (!algorithm.has_value()) {

#ifndef BE_TEST

            constexpr std::string_view msg =

                    "RowsetWriterContext::determine_encryption is not called when creating this "

                    "RowsetWriterContext, it will result in encrypted rowsets left unencrypted";

            auto st = Status::InternalError(msg);

            LOG(WARNING) << st;

            DCHECK(false) << st;

#else

            algorithm = EncryptionAlgorithmPB::PLAINTEXT;

#endif

        }

        // Apply packed file system first for write path if enabled

        // Create empty index_map for write path

        // Index information will be populated after write completes

        bool has_v1_inverted_index = tablet_schema != nullptr &&

                                     tablet_schema->has_inverted_index() &&

                                     tablet_schema->get_inverted_index_storage_format() ==

                                             InvertedIndexStorageFormatPB::V1;

        if (has_v1_inverted_index && allow_packed_file && config::enable_packed_file) {

            static constexpr std::string_view kMsg =

                    "Disable packed file for V1 inverted index tablet to avoid missing index "

                    "metadata (temporary workaround)";

            LOG(INFO) << kMsg << ", tablet_id=" << tablet_id << ", rowset_id=" << rowset_id;

        }

        // Only enable merge file for S3 file system, not for HDFS or other remote file systems

        packed_file_active = allow_packed_file && config::is_cloud_mode() &&

                             config::enable_packed_file && !has_v1_inverted_index && is_s3_fs;

        if (packed_file_active) {

            io::PackedAppendContext append_info;

            append_info.tablet_id = tablet_id;

            append_info.rowset_id = rowset_id.to_string();

            append_info.first_segment_id = first_segment_id;

            append_info.txn_id = txn_id;

            append_info.expiration_time = file_cache_ttl_sec > 0 && newest_write_timestamp > 0

                                                  ? newest_write_timestamp + file_cache_ttl_sec

                                                  : 0;

            fs = std::make_shared<io::PackedFileSystem>(fs, append_info);

        }

        // Then apply encryption on top

        if (algorithm.has_value()) {

            fs = io::make_file_system(fs, algorithm.value());

        }

        // Cache the result to ensure consistency across multiple calls

        _cached_fs = fs;

        return fs;

    }

    io::FileSystem& fs_ref() const { return *fs(); }

    io::FileWriterOptions get_file_writer_options(FileType file_type = FileType::SEGMENT_FILE) {

        io::FileWriterOptions opts {.write_file_cache = write_file_cache,

                                    .is_cold_data = is_hot_data,

                                    .file_cache_expiration_time = file_cache_ttl_sec,

                                    .approximate_bytes_to_write = approximate_bytes_to_write};

        if (config::enable_file_cache_write_index_file_only) {

            opts.allow_adaptive_file_cache_write = false;

            opts.approximate_bytes_to_write = 0;

            opts.write_file_cache = file_type == FileType::INVERTED_INDEX_FILE;

            return opts;

        }

        if (compaction_output_write_index_only && file_type == FileType::SEGMENT_FILE) {

            opts.write_file_cache = false;

            opts.allow_adaptive_file_cache_write = false;

            opts.approximate_bytes_to_write = 0;

        }

        return opts;

    }

    struct BinlogOptions {

    public:

        bool enable = false;

        void set_need_before(bool need_before) {

            this->_need_before = need_before;

            _segment_write_binlog_opt.write_before = need_before;

        }

        segment_v2::SegmentWriteBinlogOptions& write_binlog_config() {

            return _segment_write_binlog_opt;

        }

        const segment_v2::SegmentWriteBinlogOptions& write_binlog_config() const {

            return _segment_write_binlog_opt;

        }

    private:

        bool _need_before = false;

        segment_v2::SegmentWriteBinlogOptions _segment_write_binlog_opt;

    } _write_binlog_opt;

    BinlogOptions& write_binlog_opt() { return _write_binlog_opt; }

    const BinlogOptions& write_binlog_opt() const { return _write_binlog_opt; }

};

inline segment_v2::HistoricalRowRetrieverContext

RowsetWriterContext::make_historical_row_retriever_context() {

    return segment_v2::HistoricalRowRetrieverContext {

            .tablet = tablet,

            .tablet_schema = tablet_schema,

            .rowset_writer_ctx = this,

            .partial_update_info = partial_update_info,

            .is_transient_rowset_writer = is_transient_rowset_writer,

            .write_type = write_type};

}

} // namespace doris