// 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 "load/routine_load/data_consumer_group.h"

#include <gen_cpp/PlanNodes_types.h>

#include <stddef.h>

#include <map>

#include <ostream>

#include <utility>

#include "common/logging.h"

#include "io/fs/kafka_consumer_pipe.h"

#include "io/fs/kinesis_consumer_pipe.h"

#include "librdkafka/rdkafkacpp.h"

#include "load/routine_load/consumer_group_helpers.h"

#include "load/routine_load/data_consumer.h"

#include "load/stream_load/stream_load_context.h"

#include "util/stopwatch.hpp"

namespace doris {

bool DataConsumerGroup::_submit_all_consumers(

        std::function<void(std::shared_ptr<DataConsumer>, ConsumeFinishCallback)> consume_fn,

        std::function<void()> shutdown_fn, Status& result_st) {

    for (auto& consumer : _consumers) {

        auto cb = [this, shutdown_fn, &result_st](const Status& st) {

            std::unique_lock<std::mutex> lock(_mutex);

            if (--_counter == 0) {

                shutdown_fn();

                LOG(INFO) << "all consumers finished, shutdown queue. grp: " << _grp_id;

            }

            if (result_st.ok() && !st.ok()) {

                result_st = st;

            }

        };

        if (!_thread_pool.offer([consume_fn, consumer, cb] { consume_fn(consumer, cb); })) {

            LOG(WARNING) << "failed to submit consumer: " << consumer->id() << ", grp: " << _grp_id;

            return false;

        }

        VLOG_CRITICAL << "submit consumer: " << consumer->id() << ", grp: " << _grp_id;

    }

    return true;

}

Status DataConsumerGroup::_run_consume_loop(std::shared_ptr<StreamLoadContext> ctx,

                                            std::shared_ptr<io::StreamLoadPipe> pipe,

                                            Status& result_st) {

    int64_t left_time = ctx->max_interval_s * 1000;

    int64_t left_rows = ctx->max_batch_rows;

    int64_t left_bytes = ctx->max_batch_size;

    LOG(INFO) << "start consumer group: " << _grp_id << ". max time(ms): " << left_time

              << ", batch rows: " << left_rows << ", batch size: " << left_bytes << ". "

              << ctx->brief();

    MonotonicStopWatch watch;

    watch.start();

    bool eos = false;

    while (true) {

        if (eos || left_time <= 0 || left_rows <= 0 || left_bytes <= 0) {

            LOG(INFO) << "consumer group done: " << _grp_id

                      << ". consume time(ms)=" << ctx->max_interval_s * 1000 - left_time

                      << ", received rows=" << ctx->max_batch_rows - left_rows

                      << ", received bytes=" << ctx->max_batch_size - left_bytes << ", eos: " << eos

                      << ", left_time: " << left_time << ", left_rows: " << left_rows

                      << ", left_bytes: " << left_bytes

                      << ", blocking get time(us): " << pipe->get_queue_size() << ", "

                      << ctx->brief();

            _shutdown_queue();

            for (auto& consumer : _consumers) {

                static_cast<void>(consumer->cancel(ctx));

            }

            _thread_pool.shutdown();

            _thread_pool.join();

            if (!result_st.ok()) {

                pipe->cancel(result_st.to_string());

                return result_st;

            }

            RETURN_IF_ERROR(pipe->finish());

            _on_finish(ctx);

            ctx->receive_bytes = ctx->max_batch_size - left_bytes;

            return Status::OK();

        }

        if (!_dequeue_and_process(pipe.get(), left_rows, left_bytes, result_st)) {

            eos = true;

        }

        left_time = ctx->max_interval_s * 1000 - watch.elapsed_time() / 1000 / 1000;

    }

}

Status KafkaDataConsumerGroup::assign_topic_partitions(std::shared_ptr<StreamLoadContext> ctx) {

    DCHECK(ctx->kafka_info);

    DCHECK(_consumers.size() >= 1);

    // divide partitions using round-robin partitioner

    int consumer_size = doris::cast_set<int>(_consumers.size());

    auto divide_parts = WorkPartitioner<int32_t, int64_t>::partition_round_robin(

            ctx->kafka_info->begin_offset, consumer_size);

    // assign partitions to consumers equally

    for (int j = 0; j < consumer_size; ++j) {

        RETURN_IF_ERROR(

                std::static_pointer_cast<KafkaDataConsumer>(_consumers[j])

                        ->assign_topic_partitions(divide_parts[j], ctx->kafka_info->topic, ctx));

    }

    return Status::OK();

}

KafkaDataConsumerGroup::~KafkaDataConsumerGroup() {

    // clean the msgs left in queue

    _queue.shutdown();

    while (true) {

        RdKafka::Message* msg;

        if (_queue.blocking_get(&msg)) {

            delete msg;

            msg = nullptr;

        } else {

            break;

        }

    }

    DCHECK(_queue.get_size() == 0);

}

Status KafkaDataConsumerGroup::start_all(std::shared_ptr<StreamLoadContext> ctx,

                                         std::shared_ptr<io::StreamLoadPipe> pipe) {

    DORIS_CHECK(std::dynamic_pointer_cast<io::KafkaConsumerPipe>(pipe) != nullptr);

    Status result_st = Status::OK();

    _cmt_offset = ctx->kafka_info->cmt_offset;

    _format = ctx->format;

    if (!_submit_all_consumers(

                [this, max_time = ctx->max_interval_s * 1000](std::shared_ptr<DataConsumer> c,

                                                              ConsumeFinishCallback cb) {

                    actual_consume(c, &_queue, max_time, cb);

                },

                [this] { _queue.shutdown(); }, result_st)) {

        return Status::InternalError("failed to submit data consumer");

    }

    RETURN_IF_ERROR(_run_consume_loop(ctx, pipe, result_st));

    ctx->kafka_info->cmt_offset = std::move(_cmt_offset);

    return Status::OK();

}

bool KafkaDataConsumerGroup::_dequeue_and_process(io::StreamLoadPipe* pipe, int64_t& left_rows,

                                                  int64_t& left_bytes, Status& result_st) {

    RdKafka::Message* msg = nullptr;

    if (!_queue.controlled_blocking_get(&msg, config::blocking_queue_cv_wait_timeout_ms)) {

        return false;

    }

    Defer delete_msg {[msg] { delete msg; }};

    VLOG_NOTICE << "get kafka message, partition: " << msg->partition()

                << ", offset: " << msg->offset() << ", len: " << msg->len();

    if (msg->err() == RdKafka::ERR__PARTITION_EOF) {

        if (msg->offset() > 0) {

            _cmt_offset[msg->partition()] = msg->offset() - 1;

        }

        return true;

    }

    auto append_fn = FormatAppender::get_append_function<io::StreamLoadPipe>(_format);

    Status st = (pipe->*append_fn)(static_cast<const char*>(msg->payload()),

                                   static_cast<size_t>(msg->len()));

    if (st.ok()) {

        left_rows--;

        left_bytes -= msg->len();

        _cmt_offset[msg->partition()] = msg->offset();

        VLOG_NOTICE << "consume partition[" << msg->partition() << " - " << msg->offset() << "]";

    } else {

        LOG(WARNING) << "failed to append msg to pipe. grp: " << _grp_id;

        std::unique_lock<std::mutex> lock(_mutex);

        if (result_st.ok()) {

            result_st = st;

        }

    }

    return true;

}

void KafkaDataConsumerGroup::_on_finish(std::shared_ptr<StreamLoadContext> ctx) {

    // cmt_offset is moved back in start_all after _run_consume_loop returns

}

void KafkaDataConsumerGroup::actual_consume(std::shared_ptr<DataConsumer> consumer,

                                            BlockingQueue<RdKafka::Message*>* queue,

                                            int64_t max_running_time_ms, ConsumeFinishCallback cb) {

    Status st = std::static_pointer_cast<KafkaDataConsumer>(consumer)->group_consume(

            queue, max_running_time_ms);

    cb(st);

}

Status KinesisDataConsumerGroup::assign_stream_shards(std::shared_ptr<StreamLoadContext> ctx) {

    DCHECK(ctx->kinesis_info);

    DCHECK(_consumers.size() >= 1);

    // divide shards using round-robin partitioner

    int consumer_size = doris::cast_set<int>(_consumers.size());

    auto divide_shards = WorkPartitioner<std::string, std::string>::partition_round_robin(

            ctx->kinesis_info->begin_sequence_number, consumer_size);

    // assign shards to consumers equally

    for (int j = 0; j < consumer_size; ++j) {

        RETURN_IF_ERROR(std::static_pointer_cast<KinesisDataConsumer>(_consumers[j])

                                ->assign_shards(divide_shards[j], ctx->kinesis_info->stream, ctx));

    }

    return Status::OK();

}

KinesisDataConsumerGroup::~KinesisDataConsumerGroup() {

    _queue.shutdown();

    while (true) {

        std::shared_ptr<Aws::Kinesis::Model::Record> record;

        if (_queue.blocking_get(&record)) {

            record.reset();

        } else {

            break;

        }

    }

    DCHECK(_queue.get_size() == 0);

}

Status KinesisDataConsumerGroup::start_all(std::shared_ptr<StreamLoadContext> ctx,

                                           std::shared_ptr<io::StreamLoadPipe> pipe) {

    DORIS_CHECK(std::dynamic_pointer_cast<io::KinesisConsumerPipe>(pipe) != nullptr);

    Status result_st = Status::OK();

    _format = ctx->format;

    if (!_submit_all_consumers(

                [this, max_time = ctx->max_interval_s * 1000](std::shared_ptr<DataConsumer> c,

                                                              ConsumeFinishCallback cb) {

                    actual_consume(c, &_queue, max_time, cb);

                },

                [this] { _queue.shutdown(); }, result_st)) {

        return Status::InternalError("failed to submit kinesis data consumer");

    }

    return _run_consume_loop(ctx, pipe, result_st);

}

bool KinesisDataConsumerGroup::_dequeue_and_process(io::StreamLoadPipe* pipe, int64_t& left_rows,

                                                    int64_t& left_bytes, Status& result_st) {

    std::shared_ptr<Aws::Kinesis::Model::Record> record;

    if (!_queue.controlled_blocking_get(&record, config::blocking_queue_cv_wait_timeout_ms)) {

        return false;

    }

    auto& data = record->GetData();

    const char* payload = reinterpret_cast<const char*>(data.GetUnderlyingData());

    size_t len = data.GetLength();

    VLOG_NOTICE << "get kinesis record, seq: " << record->GetSequenceNumber() << ", len: " << len;

    auto append_fn = FormatAppender::get_append_function<io::StreamLoadPipe>(_format);

    Status st = (pipe->*append_fn)(payload, len);

    if (st.ok()) {

        left_rows--;

        left_bytes -= len;

        VLOG_NOTICE << "consume kinesis record [seq=" << record->GetSequenceNumber() << "]";

    } else {

        LOG(WARNING) << "failed to append kinesis record to pipe. grp: " << _grp_id;

        std::unique_lock<std::mutex> lock(_mutex);

        if (result_st.ok()) {

            result_st = st;

        }

    }

    return true;

}

void KinesisDataConsumerGroup::_on_finish(std::shared_ptr<StreamLoadContext> ctx) {

    for (auto& consumer : _consumers) {

        auto kinesis_consumer = std::static_pointer_cast<KinesisDataConsumer>(consumer);

        for (auto& [shard_id, seq_num] : kinesis_consumer->get_committed_sequence_numbers()) {

            ctx->kinesis_info->cmt_sequence_number[shard_id] = seq_num;

        }

        for (auto& [shard_id, millis] : kinesis_consumer->get_millis_behind_latest()) {

            auto [it, inserted] = ctx->kinesis_info->millis_behind_latest.emplace(shard_id, millis);

            if (!inserted && it->second < millis) {

                it->second = millis;

            }

        }

        for (auto& shard_id : kinesis_consumer->get_closed_shard_ids()) {

            ctx->kinesis_info->closed_shard_ids.insert(shard_id);

        }

    }

}

void KinesisDataConsumerGroup::actual_consume(

        std::shared_ptr<DataConsumer> consumer,

        BlockingQueue<std::shared_ptr<Aws::Kinesis::Model::Record>>* queue,

        int64_t max_running_time_ms, ConsumeFinishCallback cb) {

    Status st = std::static_pointer_cast<KinesisDataConsumer>(consumer)->group_consume(

            queue, max_running_time_ms);

    cb(st);

}

} // namespace doris