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

#include <cstdint>

#include <memory>

#include <mutex>

#include <ostream>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include <utility>

#include "common/status.h"

#include "runtime/runtime_profile.h"

#include "runtime/workload_group/workload_group_fwd.h"

#include "runtime/workload_management/resource_context.h"

#include "util/uid_util.h"

namespace doris {

class EosCompletion;

class PTabletWriterOpenRequest;

class PTabletWriterAddBlockRequest;

class PTabletWriterAddBlockResult;

class OpenPartitionRequest;

class BaseTabletsChannel;

class AtomicStatus;

// A LoadChannel manages tablets channels for all indexes

// corresponding to a certain load job

class LoadChannel {

public:

    LoadChannel(const UniqueId& load_id, int64_t timeout_s, bool is_high_priority,

                std::string sender_ip, int64_t backend_id, bool enable_profile, int64_t wg_id);

    ~LoadChannel();

    // open a new load channel if not exist

    Status open(const PTabletWriterOpenRequest& request);

    // This batch must belong to one index in one transaction. For hang_wait EOS,

    // return the sender-arrival barrier in eos_completion. Final flush/commit

    // stays synchronous on the last sender. The caller must register its RPC

    // completion only after all synchronous request/response access has ended.

    Status add_batch(const PTabletWriterAddBlockRequest& request,

                     PTabletWriterAddBlockResult* response,

                     std::shared_ptr<EosCompletion>* eos_completion);

    // return true if this load channel has been opened and all tablets channels are closed then.

    bool is_finished();

    // Publish cancellation without taking channel/writer locks. In-flight

    // requests retain ownership and release writers after their work finishes.

    Status cancel(const Status& reason = Status::Cancelled("load channel cancelled"));

    Status cancel_status() const;

    time_t last_updated_time() const { return _last_updated_time.load(); }

    const UniqueId& load_id() const { return _load_id; }

    int64_t timeout() const { return _timeout_s; }

    bool is_high_priority() const { return _is_high_priority; }

    bool is_cancelled() const;

    WorkloadGroupPtr workload_group() const { return _resource_ctx->workload_group(); }

    RuntimeProfile::Counter* get_mgr_add_batch_timer() { return _mgr_add_batch_timer; }

    RuntimeProfile::Counter* get_handle_mem_limit_timer() { return _handle_mem_limit_timer; }

protected:

    Status _get_tablets_channel(std::shared_ptr<BaseTabletsChannel>& channel, bool& is_finished,

                                int64_t index_id);

    Status _handle_eos(BaseTabletsChannel* channel, const PTabletWriterAddBlockRequest& request,

                       PTabletWriterAddBlockResult* response);

    void _init_profile();

    // thread safety

    void _report_profile(PTabletWriterAddBlockResult* response);

private:

    UniqueId _load_id;

    int64_t _txn_id = 0;

    std::mutex _profile_serialize_lock;

    std::unique_ptr<RuntimeProfile> _profile;

    RuntimeProfile* _self_profile = nullptr;

    RuntimeProfile::Counter* _add_batch_number_counter = nullptr;

    RuntimeProfile::Counter* _add_batch_timer = nullptr;

    RuntimeProfile::Counter* _add_batch_times = nullptr;

    RuntimeProfile::Counter* _mgr_add_batch_timer = nullptr;

    RuntimeProfile::Counter* _handle_mem_limit_timer = nullptr;

    RuntimeProfile::Counter* _handle_eos_timer = nullptr;

    // lock protect the tablets channel map

    std::mutex _lock;

    // index id -> tablets channel

    std::unordered_map<int64_t, std::shared_ptr<BaseTabletsChannel>> _tablets_channels;

    // index id -> (received rows, filtered rows)

    std::unordered_map<int64_t, std::pair<size_t, size_t>> _tablets_channels_rows;

    std::mutex _tablets_channels_lock;

    // This is to save finished channels id, to handle the retry request.

    std::unordered_set<int64_t> _finished_channel_ids;

    // set to true if at least one tablets channel has been opened

    bool _opened = false;

    const std::shared_ptr<AtomicStatus> _cancel_status;

    std::shared_ptr<ResourceContext> _resource_ctx;

    std::atomic<time_t> _last_updated_time;

    // the timeout of this load job.

    // Timed out channels will be periodically deleted by LoadChannelMgr.

    int64_t _timeout_s;

    // true if this is a high priority load task

    bool _is_high_priority = false;

    // the ip where tablet sink locate

    std::string _sender_ip;

    int64_t _backend_id;

    bool _enable_profile;

};

inline std::ostream& operator<<(std::ostream& os, LoadChannel& load_channel) {

    os << "LoadChannel(id=" << load_channel.load_id()

       << ", last_update_time=" << static_cast<uint64_t>(load_channel.last_updated_time())

       << ", is high priority: " << load_channel.is_high_priority() << ")";

    return os;

}

} // namespace doris