// 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 <glog/logging.h>

#include <atomic>

#include <condition_variable>

#include <cstdint>

#include <mutex>

#include <ostream>

#include <shared_mutex>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include <vector>

#include "common/status.h"

#include "core/custom_allocator.h"

#include "exec/sink/vtablet_finder.h"

#include "load/channel/adaptive_random_bucket_state.h"

#include "load/delta_writer/delta_writer_context.h"

#include "runtime/runtime_profile.h"

#include "util/bitmap.h"

#include "util/uid_util.h"

namespace google::protobuf {

template <typename Element>

class RepeatedField;

template <typename Key, typename T>

class Map;

template <typename T>

class RepeatedPtrField;

} // namespace google::protobuf

namespace doris {

class PTabletError;

class PTabletInfo;

class PTabletWriterOpenRequest;

class PTabletWriterOpenResult;

class PTabletWriterAddBlockRequest;

class PTabletWriterAddBlockResult;

class PUniqueId;

class TupleDescriptor;

class OpenPartitionRequest;

class StorageEngine;

struct TabletsChannelKey {

    UniqueId id;

    int64_t index_id;

    TabletsChannelKey(const PUniqueId& pid, int64_t index_id_) : id(pid), index_id(index_id_) {}

    ~TabletsChannelKey() noexcept = default;

    bool operator==(const TabletsChannelKey& rhs) const noexcept {

        return index_id == rhs.index_id && id == rhs.id;

    }

    std::string to_string() const;

};

std::ostream& operator<<(std::ostream& os, const TabletsChannelKey& key);

class BaseDeltaWriter;

class MemTableWriter;

class OlapTableSchemaParam;

class LoadChannel;

class EosCompletion;

struct WriteRequest;

// Write channel for a particular (load, index).

class BaseTabletsChannel {

public:

    BaseTabletsChannel(const TabletsChannelKey& key, const UniqueId& load_id, bool is_high_priority,

                       RuntimeProfile* profile);

    virtual ~BaseTabletsChannel();

    Status open(const PTabletWriterOpenRequest& request);

    // open + open writers

    Status incremental_open(const PTabletWriterOpenRequest& params);

    virtual std::unique_ptr<BaseDeltaWriter> create_delta_writer(const WriteRequest& request) = 0;

    // no-op when this channel has been closed or cancelled

    virtual Status add_batch(const PTabletWriterAddBlockRequest& request,

                             PTabletWriterAddBlockResult* response) = 0;

    // Mark sender with 'sender_id' as closed.

    // If all senders are closed, close this channel, set '*finished' to true, update 'tablet_vec'

    // to include all tablets written in this channel.

    // Release the EOS arrival barrier before final flushing/committing, outside _lock.

    // no-op when this channel has been closed or cancelled

    virtual Status close(LoadChannel* parent, const PTabletWriterAddBlockRequest& req,

                         PTabletWriterAddBlockResult* res, bool* finished) = 0;

    // no-op when this channel has been closed or cancelled

    virtual Status cancel();

    const std::shared_ptr<EosCompletion>& eos_completion() const { return _eos_completion; }

    void refresh_profile();

    size_t total_received_rows() const { return _total_received_rows; }

    size_t num_rows_filtered() const { return _num_rows_filtered; }

    // means this tablets in this BE is incremental opened partitions.

    bool is_incremental_channel() const { return _open_by_incremental; }

    bool is_finished() const { return _state == kFinished; }

protected:

    // Called by the final sender under _lock. Publish arrival before flushing,

    // releasing the channel lock while invoking RPC callbacks.

    Status _close_and_notify(std::unique_lock<std::mutex>& lock, int sender_id,

                             PTabletWriterAddBlockResult* response);

    Status _get_close_result(std::unique_lock<std::mutex>& lock, int sender_id,

                             PTabletWriterAddBlockResult* response);

    // Called exactly once by the final sender, with _lock held.

    virtual Status _close_writers(PTabletWriterAddBlockResult* response) = 0;

    Status _init_adaptive_random_bucket_state(const PTabletWriterOpenRequest& request);

    Status _write_block_data(const PTabletWriterAddBlockRequest& request, int64_t cur_seq,

                             std::unordered_map<int64_t, TabletAddRowsPayload>& tablet_to_rows,

                             PTabletWriterAddBlockResult* response);

    Status _write_block_data_for_adaptive_random_bucket(

            const PTabletWriterAddBlockRequest& request, int64_t cur_seq,

            std::unordered_map<int64_t, DorisVector<uint32_t>>& partition_to_rowidxs,

            PTabletWriterAddBlockResult* response);

    virtual Status _prepare_adaptive_random_bucket_writer(BaseDeltaWriter* writer);

    Status _build_partition_to_rowidxs_for_adaptive_random_bucket(

            const PTabletWriterAddBlockRequest& request,

            std::unordered_map<int64_t, DorisVector<uint32_t>>* partition_to_rowidxs);

    std::shared_ptr<std::mutex> _get_partition_route_lock(int64_t partition_id);

    // OK does not imply a write: duplicates and closed channels are no-ops.

    Status _get_current_seq(int64_t& cur_seq, const PTabletWriterAddBlockRequest& request,

                            bool& should_write);

    // open all writer

    Status _open_all_writers(const PTabletWriterOpenRequest& request);

    void _add_broken_tablet(int64_t tablet_id);

    // thread-unsafe, add a shared lock for `_tablet_writers_lock` if needed

    bool _is_broken_tablet(int64_t tablet_id) const;

    void _add_error_tablet(google::protobuf::RepeatedPtrField<PTabletError>* tablet_errors,

                           int64_t tablet_id, Status error) const;

    void _build_tablet_to_rows(

            const PTabletWriterAddBlockRequest& request,

            std::unordered_map<int64_t /* tablet_id */, TabletAddRowsPayload>* tablet_to_rows);

    virtual void _init_profile(RuntimeProfile* profile);

    // id of this load channel

    TabletsChannelKey _key;

    std::shared_ptr<EosCompletion> _eos_completion;

    // protect _state change. open and close. when add_batch finished, lock to change _next_seqs also

    std::mutex _lock;

    enum State {

        kInitialized,

        kOpened,

        kFinished // closed or cancelled

    };

    State _state;

    UniqueId _load_id;

    // initialized in open function

    int64_t _txn_id = -1;

    int64_t _index_id = -1;

    std::shared_ptr<OlapTableSchemaParam> _schema;

    TupleDescriptor* _tuple_desc = nullptr;

    bool _open_by_incremental = false;

    // next sequence we expect

    std::set<int32_t> _recieved_senders;

    int _num_remaining_senders = 0;

    std::vector<int64_t> _next_seqs;

    Bitmap _closed_senders;

    // Arrival releases earlier senders; only the final sender's retries must

    // wait for flush/commit and receive its status and tablet results. All these

    // fields are protected by _lock, including while callbacks run without it.

    int _final_sender_id = -1;

    bool _final_close_in_progress = false;

    std::condition_variable _final_close_cv;

    Status _close_status;

    std::unique_ptr<PTabletWriterAddBlockResult> _final_close_result;

    // tablet_id -> TabletChannel. it will only be changed in open() or inc_open()

    std::unordered_map<int64_t, std::unique_ptr<BaseDeltaWriter>> _tablet_writers;

    // protect _tablet_writers

    std::mutex _tablet_writers_lock;

    // broken tablet ids.

    // If a tablet write fails, it's id will be added to this set.

    // So that following batch will not handle this tablet anymore.

    std::unordered_set<int64_t> _broken_tablets;

    std::shared_mutex _broken_tablets_lock;

    std::unordered_set<int64_t> _reducing_tablets;

    std::unordered_set<int64_t> _partition_ids;

    std::shared_ptr<AdaptiveRandomBucketState> _adaptive_random_bucket_state;

    // Protects the route-lock map. Each entry serializes current-tablet selection, write,

    // and rotation for one partition so all senders on this BE share one current bucket.

    std::mutex _partition_route_locks_lock;

    std::unordered_map<int64_t, std::shared_ptr<std::mutex>> _partition_route_locks;

    static std::atomic<uint64_t> _s_tablet_writer_count;

    bool _is_high_priority = false;

    RuntimeProfile* _profile = nullptr;

    RuntimeProfile::Counter* _add_batch_number_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _memory_usage_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _write_memory_usage_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _flush_memory_usage_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _max_tablet_memory_usage_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _max_tablet_write_memory_usage_counter = nullptr;

    RuntimeProfile::HighWaterMarkCounter* _max_tablet_flush_memory_usage_counter = nullptr;

    RuntimeProfile::Counter* _add_batch_timer = nullptr;

    RuntimeProfile::Counter* _write_block_timer = nullptr;

    RuntimeProfile::Counter* _incremental_open_timer = nullptr;

    // record rows received and filtered

    size_t _total_received_rows = 0;

    size_t _num_rows_filtered = 0;

};

class DeltaWriter;

// `StorageEngine` mixin for `BaseTabletsChannel`

class TabletsChannel final : public BaseTabletsChannel {

public:

    TabletsChannel(StorageEngine& engine, const TabletsChannelKey& key, const UniqueId& load_id,

                   bool is_high_priority, RuntimeProfile* profile);

    ~TabletsChannel() override;

    std::unique_ptr<BaseDeltaWriter> create_delta_writer(const WriteRequest& request) override;

    Status add_batch(const PTabletWriterAddBlockRequest& request,

                     PTabletWriterAddBlockResult* response) override;

    Status close(LoadChannel* parent, const PTabletWriterAddBlockRequest& req,

                 PTabletWriterAddBlockResult* res, bool* finished) override;

private:

    Status _close_writers(PTabletWriterAddBlockResult* response) override;

    // deal with DeltaWriter commit_txn(), add tablet to list for return.

    void _commit_txn(DeltaWriter* writer, PTabletWriterAddBlockResult* res);

    StorageEngine& _engine;

};

} // namespace doris