// Licensed to the LF AI & Data foundation 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 "Driver.h"

#include <folly/ExceptionWrapper.h>
#include <folly/Try.h>
#include <algorithm>
#include <atomic>
#include <cassert>
#include <exception>
#include <memory>
#include <string>

#include "common/Common.h"
#include "common/EasyAssert.h"
#include "common/Exception.h"
#include "common/Tracer.h"
#include "common/protobuf_utils.h"
#include "exec/QueryContext.h"
#include "exec/Task.h"
#include "exec/operator/AggregationNode.h"
#include "exec/operator/CallbackSink.h"
#include "exec/operator/ElementFilterBitsNode.h"
#include "exec/operator/IterativeElementFilterNode.h"
#include "exec/operator/FilterBitsNode.h"
#include "exec/operator/IterativeFilterNode.h"
#include "exec/operator/MvccNode.h"
#include "exec/operator/Operator.h"
#include "exec/operator/ProjectNode.h"
#include "exec/operator/RandomSampleNode.h"
#include "exec/operator/RescoresNode.h"
#include "exec/operator/SearchGroupByNode.h"
#include "exec/operator/VectorSearchNode.h"
#include "exec/operator/QueryOrderByNode.h"
#include "fmt/core.h"
#include "folly/Executor.h"
#include "folly/Unit.h"
#include "glog/logging.h"
#include "log/Log.h"
#include "plan/PlanNode.h"
#include "storage/PrefetchThreadPool.h"

namespace milvus {
namespace exec {

std::atomic_uint64_t BlockingState::num_blocked_drivers_{0};

std::shared_ptr<QueryConfig>
DriverContext::GetQueryConfig() {
    return task_->query_context()->query_config();
}

std::shared_ptr<Driver>
DriverFactory::CreateDriver(
    std::unique_ptr<DriverContext> ctx,
    const std::function<int(int pipelineid)>& num_drivers) {
    auto driver = std::shared_ptr<Driver>(new Driver());
    ctx->driver_ = driver.get();
    std::vector<std::shared_ptr<Operator>> operators;
    operators.reserve(plannodes_.size());

    for (size_t i = 0; i < plannodes_.size(); ++i) {
        auto id = operators.size();
        const auto& plannode = plannodes_[i];
        if (auto filterbitsnode =
                std::dynamic_pointer_cast<const plan::FilterBitsNode>(
                    plannode)) {
            tracer::AddEvent("create_operator: FilterBitsNode");
            operators.push_back(std::make_shared<PhyFilterBitsNode>(
                id, ctx.get(), filterbitsnode));
        } else if (auto filternode = std::dynamic_pointer_cast<
                       const plan::IterativeFilterNode>(plannode)) {
            tracer::AddEvent("create_operator: IterativeFilterNode");
            operators.push_back(std::make_shared<PhyIterativeFilterNode>(
                id, ctx.get(), filternode));
        } else if (auto mvccnode =
                       std::dynamic_pointer_cast<const plan::MvccNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: MvccNode");
            operators.push_back(
                std::make_shared<PhyMvccNode>(id, ctx.get(), mvccnode));
        } else if (auto vectorsearchnode =
                       std::dynamic_pointer_cast<const plan::VectorSearchNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: VectorSearchNode");
            operators.push_back(std::make_shared<PhyVectorSearchNode>(
                id, ctx.get(), vectorsearchnode));
        } else if (auto searchGroupByNode =
                       std::dynamic_pointer_cast<const plan::SearchGroupByNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: SearchGroupByNode");
            operators.push_back(std::make_shared<PhySearchGroupByNode>(
                id, ctx.get(), searchGroupByNode));
        } else if (auto queryGroupByNode =
                       std::dynamic_pointer_cast<const plan::AggregationNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: AggregationNode");
            operators.push_back(std::make_shared<PhyAggregationNode>(
                id, ctx.get(), queryGroupByNode));
        } else if (auto projectNode =
                       std::dynamic_pointer_cast<const plan::ProjectNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: ProjectNode");
            operators.push_back(
                std::make_shared<PhyProjectNode>(id, ctx.get(), projectNode));
        } else if (auto orderByNode =
                       std::dynamic_pointer_cast<const plan::OrderByNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: QueryOrderByNode");
            operators.push_back(std::make_shared<PhyQueryOrderByNode>(
                id, ctx.get(), orderByNode));
        } else if (auto samplenode =
                       std::dynamic_pointer_cast<const plan::RandomSampleNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: RandomSampleNode");
            operators.push_back(std::make_shared<PhyRandomSampleNode>(
                id, ctx.get(), samplenode));
        } else if (auto rescoresnode =
                       std::dynamic_pointer_cast<const plan::RescoresNode>(
                           plannode)) {
            tracer::AddEvent("create_operator: RescoresNode");
            operators.push_back(
                std::make_shared<PhyRescoresNode>(id, ctx.get(), rescoresnode));
        } else if (auto node = std::dynamic_pointer_cast<
                       const plan::IterativeElementFilterNode>(plannode)) {
            tracer::AddEvent("create_operator: IterativeElementFilterNode");
            operators.push_back(std::make_shared<PhyIterativeElementFilterNode>(
                id, ctx.get(), node));
        } else if (auto node = std::dynamic_pointer_cast<
                       const plan::ElementFilterBitsNode>(plannode)) {
            tracer::AddEvent("create_operator: ElementFilterBitsNode");
            operators.push_back(std::make_shared<PhyElementFilterBitsNode>(
                id, ctx.get(), node));
        } else {
            ThrowInfo(ErrorCode::UnexpectedError, "Unknown plan node type");
        }
        // TODO: add more operators
    }

    if (consumer_supplier_) {
        operators.push_back(consumer_supplier_(operators.size(), ctx.get()));
    }

    driver->Init(std::move(ctx), std::move(operators));

    return driver;
}

void
Driver::Enqueue(std::shared_ptr<Driver> driver) {
    if (driver->closed_) {
        return;
    }

    driver->get_task()->query_context()->executor()->add(
        [driver]() { Driver::Run(driver); });
}

void
Driver::Run(std::shared_ptr<Driver> self) {
    std::shared_ptr<BlockingState> blocking_state;
    RowVectorPtr result;
    auto reason = self->RunInternal(self, blocking_state, result);

    AssertInfo(result == nullptr,
               "The last operator (sink) must not produce any results.");

    if (reason == StopReason::kBlock) {
        return;
    }

    switch (reason) {
        case StopReason::kBlock:
            BlockingState::SetResume(blocking_state);
            return;
        case StopReason::kYield:
            Enqueue(self);
        case StopReason::kPause:
        case StopReason::kTerminate:
        case StopReason::kAlreadyTerminated:
        case StopReason::kAtEnd:
            return;
        default:
            AssertInfo(false, "Unhandled stop reason");
    }
}

void
Driver::initializeOperators() {
    // Atomically check and set: only the first thread to call this will
    // get false (the previous value) and proceed with initialization.
    // Other threads will get true and skip initialization.
    // Use memory barriers to ensure initialization writes are visible
    // before the flag is seen by other threads.
    if (!operatorsInitialized_.exchange(true, std::memory_order_acq_rel)) {
        // This thread won the initialization race. Perform initialization.
        for (auto& op : operators_) {
            op->initialize();
        }
        // Use release semantics to ensure all initialization writes are
        // visible to other threads before they see the flag as true.
        // The flag was already set to true by exchange above, but this
        // barrier ensures all writes from initialization are visible.
        std::atomic_thread_fence(std::memory_order_release);
    } else {
        // Another thread is initializing or has completed initialization.
        // Use acquire semantics to ensure we see all initialization writes
        // after the flag becomes true.
        operatorsInitialized_.load(std::memory_order_acquire);
    }
}

void
Driver::Init(std::unique_ptr<DriverContext> ctx,
             std::vector<std::shared_ptr<Operator>> operators) {
    assert(ctx != nullptr);
    ctx_ = std::move(ctx);
    AssertInfo(operators.size() != 0, "operators in driver must not empty");
    operators_ = std::move(operators);
    current_operator_index_ = operators_.size() - 1;
}

void
Driver::Close() {
    if (closed_) {
        return;
    }

    for (auto& op : operators_) {
        op->WaitPrefetch();
        op->Close();
    }

    closed_ = true;

    Task::RemoveDriver(ctx_->task_, this);
}

RowVectorPtr
Driver::Next(std::shared_ptr<BlockingState>& blocking_state) {
    auto self = shared_from_this();

    RowVectorPtr result;
    auto stop = RunInternal(self, blocking_state, result);

    Assert(stop == StopReason::kBlock || stop == StopReason::kAtEnd ||
           stop == StopReason::kAlreadyTerminated);
    return result;
}

#define CALL_OPERATOR(call_func, operator, method_name)            \
    try {                                                          \
        call_func;                                                 \
    } catch (std::exception & e) {                                 \
        std::string stack_trace = milvus::impl::EasyStackTrace();  \
        auto err_msg = fmt::format(                                \
            "Operator::{} failed for [Operator:{}, plan node id: " \
            "{}] : {}\nStack trace: {}",                           \
            method_name,                                           \
            operator->ToString(),                                  \
            operator->get_plannode_id(),                           \
            e.what(),                                              \
            stack_trace);                                          \
        LOG_ERROR("{}", err_msg);                                  \
        throw ExecOperatorException(err_msg);                      \
    }

StopReason
Driver::RunInternal(std::shared_ptr<Driver>& self,
                    std::shared_ptr<BlockingState>& blocking_state,
                    RowVectorPtr& result) {
    try {
        initializeOperators();
        if (ENABLE_DRIVER_PREFETCH.load(std::memory_order_relaxed)) {
            std::call_once(self->once_, [self]() { self->PrefetchAsync(); });
        }
        int num_operators = operators_.size();
        ContinueFuture future;

        for (;;) {
            for (int32_t i = num_operators - 1; i >= 0; --i) {
                auto op = operators_[i].get();
                current_operator_index_ = i;
                CALL_OPERATOR(
                    blocking_reason_ = op->IsBlocked(&future), op, "IsBlocked");
                if (blocking_reason_ != BlockingReason::kNotBlocked) {
                    blocking_state = std::make_shared<BlockingState>(
                        self, std::move(future), op, blocking_reason_);
                    return StopReason::kBlock;
                }
                Operator* next_op = nullptr;

                if (i < operators_.size() - 1) {
                    next_op = operators_[i + 1].get();
                    CALL_OPERATOR(
                        blocking_reason_ = next_op->IsBlocked(&future),
                        next_op,
                        "IsBlocked");
                    if (blocking_reason_ != BlockingReason::kNotBlocked) {
                        blocking_state = std::make_shared<BlockingState>(
                            self, std::move(future), next_op, blocking_reason_);
                        return StopReason::kBlock;
                    }

                    bool needs_input;
                    CALL_OPERATOR(needs_input = next_op->NeedInput(),
                                  next_op,
                                  "NeedInput");
                    if (needs_input) {
                        RowVectorPtr result;
                        {
                            CALL_OPERATOR(
                                result = op->GetOutput(), op, "GetOutput");
                            if (result) {
                                AssertInfo(
                                    result->size() > 0,
                                    fmt::format(
                                        "GetOutput must return nullptr or "
                                        "a non-empty vector: {}",
                                        op->get_operator_type()));
                            }
                        }
                        if (result) {
                            CALL_OPERATOR(
                                next_op->AddInput(result), next_op, "AddInput");
                            i += 2;
                            continue;
                        } else {
                            CALL_OPERATOR(
                                blocking_reason_ = op->IsBlocked(&future),
                                op,
                                "IsBlocked");
                            if (blocking_reason_ !=
                                BlockingReason::kNotBlocked) {
                                blocking_state =
                                    std::make_shared<BlockingState>(
                                        self,
                                        std::move(future),
                                        next_op,
                                        blocking_reason_);
                                return StopReason::kBlock;
                            }
                            if (op->IsFinished()) {
                                CALL_OPERATOR(next_op->NoMoreInput(),
                                              next_op,
                                              "NoMoreInput");
                                break;
                            }
                        }
                    }
                } else {
                    {
                        CALL_OPERATOR(
                            result = op->GetOutput(), op, "GetOutput");
                        if (result) {
                            AssertInfo(
                                result->size() > 0,
                                fmt::format("GetOutput must return nullptr or "
                                            "a non-empty vector: {}",
                                            op->get_operator_type()));
                            blocking_reason_ = BlockingReason::kWaitForConsumer;
                            return StopReason::kBlock;
                        }
                    }
                    if (op->IsFinished()) {
                        Close();
                        return StopReason::kAtEnd;
                    }
                    continue;
                }
            }
        }
    } catch (std::exception& e) {
        get_task()->SetError(std::current_exception());
        return StopReason::kAlreadyTerminated;
    }
}

void
Driver::PrefetchAsync() {
    auto prefetch_pool = GetPrefetchThreadPool();
    for (auto& op : operators_) {
        op->PrefetchAsync(prefetch_pool);
    }
}

static bool
MustStartNewPipeline(std::shared_ptr<const plan::PlanNode> plannode,
                     int source_id) {
    //TODO: support LocalMerge and other shuffle
    return source_id != 0;
}

OperatorSupplier
MakeConsumerSupplier(ConsumerSupplier supplier) {
    if (supplier) {
        return [supplier](int32_t operator_id, DriverContext* ctx) {
            return std::make_unique<CallbackSink>(operator_id, ctx, supplier());
        };
    }
    return nullptr;
}

uint32_t
MaxDrivers(const DriverFactory* factory, const QueryConfig& config) {
    return 1;
}

static void
SplitPlan(const std::shared_ptr<const plan::PlanNode>& plannode,
          std::vector<std::shared_ptr<const plan::PlanNode>>* current_plannodes,
          const std::shared_ptr<const plan::PlanNode>& consumer_node,
          OperatorSupplier operator_supplier,
          std::vector<std::unique_ptr<DriverFactory>>* driver_factories) {
    if (!current_plannodes) {
        driver_factories->push_back(std::make_unique<DriverFactory>());
        current_plannodes = &driver_factories->back()->plannodes_;
        driver_factories->back()->consumer_supplier_ = operator_supplier;
        driver_factories->back()->consumer_node_ = consumer_node;
    }

    auto sources = plannode->sources();
    if (sources.empty()) {
        driver_factories->back()->is_input_driver_ = true;
    } else {
        for (int i = 0; i < sources.size(); ++i) {
            SplitPlan(
                sources[i],
                MustStartNewPipeline(plannode, i) ? nullptr : current_plannodes,
                plannode,
                nullptr,
                driver_factories);
        }
    }
    current_plannodes->push_back(plannode);
}

void
LocalPlanner::Plan(
    const plan::PlanFragment& fragment,
    ConsumerSupplier consumer_supplier,
    std::vector<std::unique_ptr<DriverFactory>>* driver_factories,
    const QueryConfig& config,
    uint32_t max_drivers) {
    SplitPlan(fragment.plan_node_,
              nullptr,
              nullptr,
              MakeConsumerSupplier(consumer_supplier),
              driver_factories);

    (*driver_factories)[0]->is_output_driver_ = true;

    for (auto& factory : *driver_factories) {
        factory->max_drivers_ = MaxDrivers(factory.get(), config);
        factory->num_drivers_ = std::min(factory->max_drivers_, max_drivers);

        if (factory->is_group_execution_) {
            factory->num_total_drivers_ =
                factory->num_drivers_ * fragment.num_splitgroups_;
        } else {
            factory->num_total_drivers_ = factory->num_drivers_;
        }
    }
}

}  // namespace exec
}  // namespace milvus