// 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 "RescoresNode.h"

#include <algorithm>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <optional>
#include <ratio>
#include <utility>

#include "common/EasyAssert.h"
#include "common/QueryResult.h"
#include "common/Tracer.h"
#include "common/Types.h"
#include "common/Utils.h"
#include "exec/QueryContext.h"
#include "exec/expression/EvalCtx.h"
#include "exec/expression/Expr.h"
#include "exec/operator/Utils.h"
#include "expr/ITypeExpr.h"
#include "fmt/core.h"
#include "folly/FBVector.h"
#include "knowhere/comp/index_param.h"
#include "monitor/Monitor.h"
#include "pb/plan.pb.h"
#include "plan/PlanNode.h"
#include "prometheus/histogram.h"
#include "rescores/BoostScoreRunner.h"
#include "rescores/Scorer.h"

namespace milvus::exec {

PhyRescoresNode::PhyRescoresNode(
    int32_t operator_id,
    DriverContext* ctx,
    const std::shared_ptr<const plan::RescoresNode>& scorer)
    : Operator(ctx,
               scorer->output_type(),
               operator_id,
               scorer->id(),
               "PhyRescoresNode") {
    scorers_ = scorer->scorers();
    option_ = scorer->option();
};

void
PhyRescoresNode::AddInput(RowVectorPtr& input) {
    input_ = std::move(input);
}

bool
PhyRescoresNode::IsFinished() {
    return is_finished_;
}

RowVectorPtr
PhyRescoresNode::GetOutput() {
    ExecContext* exec_context = operator_context_->get_exec_context();
    auto query_context_ = exec_context->get_query_context();
    milvus::exec::checkCancellation(query_context_);

    if (is_finished_ || !no_more_input_) {
        return nullptr;
    }

    tracer::AutoSpan span(
        "PhyRescoresNode::Execute", tracer::GetRootSpan(), true);

    DeferLambda([&]() { is_finished_ = true; });

    if (input_ == nullptr) {
        return nullptr;
    }

    std::chrono::high_resolution_clock::time_point scalar_start =
        std::chrono::high_resolution_clock::now();

    auto query_info = exec_context->get_query_config();
    milvus::SearchResult search_result = query_context_->get_search_result();
    auto segment = query_context_->get_segment();
    auto op_context = query_context_->get_op_context();

    // prepare segment offset
    FixedVector<int32_t> offsets;
    std::vector<size_t> offset_idx;
    offsets.reserve(search_result.seg_offsets_.size());
    offset_idx.reserve(search_result.seg_offsets_.size());

    for (size_t i = 0; i < search_result.seg_offsets_.size(); i++) {
        // remain offset will be placeholder(-1) if result count not enough (less than topk)
        // skip placeholder offset
        if (search_result.seg_offsets_[i] >= 0) {
            offsets.push_back(
                static_cast<int32_t>(search_result.seg_offsets_[i]));
            offset_idx.push_back(i);
        }
    }

    // skip rescore if result was empty
    if (offsets.empty()) {
        query_context_->set_search_result(std::move(search_result));
        return input_;
    }

    std::vector<std::optional<float>> boost_scores(offsets.size());
    auto function_mode = option_->function_mode();
    rescores::ComputeFunctionScores(exec_context,
                                    op_context,
                                    segment,
                                    scorers_,
                                    function_mode,
                                    offsets,
                                    boost_scores);

    // calculate final score
    auto boost_mode = option_->boost_mode();
    switch (boost_mode) {
        case proto::plan::BoostModeMultiply:
            for (auto i = 0; i < offsets.size(); i++) {
                if (boost_scores[i].has_value()) {
                    search_result.distances_[offset_idx[i]] *=
                        boost_scores[i].value();
                }
            }
            break;
        case proto::plan::BoostModeSum:
            for (auto i = 0; i < offsets.size(); i++) {
                if (boost_scores[i].has_value()) {
                    search_result.distances_[offset_idx[i]] +=
                        boost_scores[i].value();
                }
            }

            break;
        default:
            ThrowInfo(ErrorCode::UnexpectedError,
                      fmt::format("unknown boost boost mode: {}",
                                  static_cast<int>(boost_mode)));
    }

    knowhere::MetricType metric_type = query_context_->get_metric_type();
    bool large_is_better = PositivelyRelated(metric_type);
    sort_search_result(search_result, large_is_better);
    query_context_->set_search_result(std::move(search_result));

    std::chrono::high_resolution_clock::time_point scalar_end =
        std::chrono::high_resolution_clock::now();
    double scalar_cost =
        std::chrono::duration<double, std::micro>(scalar_end - scalar_start)
            .count();
    milvus::monitor::internal_core_search_latency_rescore.Observe(scalar_cost /
                                                                  1000);

    tracer::AddEvent(fmt::format("rescored_count: {}", offsets.size()));

    return input_;
};

}  // namespace milvus::exec