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

#include <memory>

#include <string_view>

#include <utility>

#include <vector>

#include "common/cast_set.h"

#include "common/status.h"

#include "pipeline/exec/operator.h"

#include "util/runtime_profile.h"

namespace doris::pipeline {

#include "common/compile_check_begin.h"

class PipelineFragmentContext;

class Pipeline;

using PipelinePtr = std::shared_ptr<Pipeline>;

using Pipelines = std::vector<PipelinePtr>;

using PipelineId = uint32_t;

class Pipeline : public std::enable_shared_from_this<Pipeline> {

    friend class PipelineTask;

    friend class PipelineFragmentContext;

public:

    explicit Pipeline(PipelineId pipeline_id, int num_tasks, int num_tasks_of_parent)

            : _pipeline_id(pipeline_id),

              _num_tasks(num_tasks),

              _num_tasks_of_parent(num_tasks_of_parent) {

        _init_profile();

        _tasks.resize(_num_tasks, nullptr);

    }

    // Add operators for pipelineX

    Status add_operator(OperatorPtr& op, const int parallelism);

    // prepare operators for pipelineX

    Status prepare(RuntimeState* state);

    Status set_sink(DataSinkOperatorPtr& sink_operator);

    DataSinkOperatorXBase* sink() { return _sink.get(); }

    Operators& operators() { return _operators; }

    DataSinkOperatorPtr sink_shared_pointer() { return _sink; }

    [[nodiscard]] const RowDescriptor& output_row_desc() const {

        return _operators.back()->row_desc();

    }

    [[nodiscard]] PipelineId id() const { return _pipeline_id; }

    static bool is_hash_exchange(TLocalPartitionType::type idx) {

        return is_shuffled_exchange(idx);

    }

    // For HASH_SHUFFLE, BUCKET_HASH_SHUFFLE, and ADAPTIVE_PASSTHROUGH,

    // data is processed and shuffled on the sink.

    // Compared to PASSTHROUGH, this is a relatively heavy operation.

    static bool heavy_operations_on_the_sink(TLocalPartitionType::type idx) {

        return is_shuffled_exchange(idx) || idx == TLocalPartitionType::ADAPTIVE_PASSTHROUGH;

    }

    bool need_to_local_exchange(const DataDistribution target_data_distribution,

                                const int idx) const;

    void init_data_distribution(RuntimeState* state) {

        set_data_distribution(_operators.front()->required_data_distribution(state));

    }

    void set_data_distribution(const DataDistribution& data_distribution) {

        _data_distribution = data_distribution;

    }

    const DataDistribution& data_distribution() const { return _data_distribution; }

    std::vector<std::shared_ptr<Pipeline>>& children() { return _children; }

    void set_children(std::shared_ptr<Pipeline> child) { _children.push_back(child); }

    void set_children(std::vector<std::shared_ptr<Pipeline>> children) {

        _children = std::move(children);

    }

    void incr_created_tasks(int i, PipelineTask* task) {

        _num_tasks_created++;

        _num_tasks_running++;

        DCHECK_LT(i, _tasks.size());

        _tasks[i] = task;

    }

    void make_all_runnable(PipelineId wake_by);

    void set_num_tasks(int num_tasks) {

        _num_tasks = num_tasks;

        _tasks.resize(_num_tasks, nullptr);

        for (auto& op : _operators) {

            op->set_parallel_tasks(_num_tasks);

        }

#ifndef NDEBUG

        if (num_tasks > 1 &&

            std::any_of(_operators.begin(), _operators.end(),

                        [&](OperatorPtr op) -> bool { return op->is_serial_operator(); })) {

            DCHECK(false) << debug_string();

        }

#endif

    }

    int num_tasks() const { return _num_tasks; }

    bool close_task() { return _num_tasks_running.fetch_sub(1) == 1; }

    std::string debug_string() const {

        fmt::memory_buffer debug_string_buffer;

        fmt::format_to(debug_string_buffer,

                       "Pipeline [id: {}, _num_tasks: {}, _num_tasks_created: {}]", _pipeline_id,

                       _num_tasks, _num_tasks_created);

        for (int i = 0; i < _operators.size(); i++) {

            fmt::format_to(debug_string_buffer, "\n{}", _operators[i]->debug_string(i));

        }

        fmt::format_to(debug_string_buffer, "\n{}",

                       _sink ? _sink->debug_string(cast_set<int>(_operators.size())) : "null");

        return fmt::to_string(debug_string_buffer);

    }

    int num_tasks_of_parent() const { return _num_tasks_of_parent; }

    std::string& name() { return _name; }

private:

    void _init_profile();

    std::vector<std::pair<int, std::weak_ptr<Pipeline>>> _parents;

    std::vector<std::pair<int, std::shared_ptr<Pipeline>>> _dependencies;

    std::vector<std::shared_ptr<Pipeline>> _children;

    PipelineId _pipeline_id;

    // pipline id + operator names. init when:

    //  build_operators(), if pipeline;

    //  _build_pipelines() and _create_tree_helper(), if pipelineX.

    std::string _name;

    std::unique_ptr<RuntimeProfile> _pipeline_profile;

    // Operators for pipelineX. All pipeline tasks share operators from this.

    // [SourceOperator -> ... -> SinkOperator]

    Operators _operators;

    DataSinkOperatorPtr _sink = nullptr;

    std::shared_ptr<ObjectPool> _obj_pool;

    // Input data distribution of this pipeline. We do local exchange when input data distribution

    // does not match the target data distribution.

    DataDistribution _data_distribution {TLocalPartitionType::NOOP};

    // How many tasks should be created ?

    int _num_tasks = 1;

    // How many tasks are already created?

    std::atomic<int> _num_tasks_created = 0;

    // How many tasks are already created and not finished?

    std::atomic<int> _num_tasks_running = 0;

    // Tasks in this pipeline.

    std::vector<PipelineTask*> _tasks;

    // Parallelism of parent pipeline.

    const int _num_tasks_of_parent;

};

#include "common/compile_check_end.h"

} // namespace doris::pipeline