// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed 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 <stdlib.h>
#include <common/EasyAssert.h>
#include <folly/CancellationToken.h>
#include <folly/futures/Future.h>
#include <folly/futures/SharedPromise.h>
#include "future_c_types.h"
#include "LeakyResult.h"
#include "Ready.h"
#include "pb/cgo_msg.pb.h"
#include <chrono>
#include "monitor/Monitor.h"

namespace milvus::futures {

enum class PoolType { kSearch, kLoad };

template <class Duration>
class Metrics;

template <class Duration>
class Metrics {
 public:
    class ExecutionGuard {
     public:
        explicit ExecutionGuard(Metrics& metrics) : metrics_(metrics) {
            metrics.startExecute();
        }

        ExecutionGuard(const ExecutionGuard&) = delete;
        ExecutionGuard(const ExecutionGuard&&) = delete;
        ExecutionGuard&
        operator=(const ExecutionGuard&) = delete;
        ExecutionGuard&
        operator=(const ExecutionGuard&&) = delete;

        ~ExecutionGuard() {
            metrics_.executeDone();
        }

     private:
        Metrics& metrics_;
    };

    explicit Metrics(PoolType pool_type)
        : pool_type_(pool_type),
          time_point_(std::chrono::steady_clock::now()),
          queue_duration_(0),
          execute_duration_(0) {
        inflightGauge().Increment();
    }

    Metrics(const Metrics&) = delete;
    Metrics(const Metrics&&) = delete;
    Metrics&
    operator=(const Metrics&) = delete;
    Metrics&
    operator=(const Metrics&&) = delete;

    ~Metrics() {
        inflightGauge().Decrement();
        queueDurationHist().Observe(
            std::chrono::duration<double>(queue_duration_).count());
        if (cancelled_before_execute_) {
            cancelBeforeCounter().Increment();
        } else {
            if (cancelled_during_execute_) {
                cancelDuringCounter().Increment();
            }
            executeDurationHist().Observe(
                std::chrono::duration<double>(execute_duration_).count());
        }
    }

    void
    withEarlyCancel() {
        queue_duration_ = std::chrono::duration_cast<Duration>(
            std::chrono::steady_clock::now() - time_point_);
        cancelled_before_execute_ = true;
    }

    void
    startExecute() {
        auto now = std::chrono::steady_clock::now();
        queue_duration_ =
            std::chrono::duration_cast<Duration>(now - time_point_);
        time_point_ = now;
        executingGauge().Increment();
    }

    void
    withDuringCancel() {
        cancelled_during_execute_ = true;
    }

    void
    executeDone() {
        auto now = std::chrono::steady_clock::now();
        execute_duration_ =
            std::chrono::duration_cast<Duration>(now - time_point_);
        executingGauge().Decrement();
    }

 private:
    prometheus::Gauge&
    inflightGauge() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::internal_cgo_inflight_task_total_search
                   : milvus::monitor::internal_cgo_inflight_task_total_load;
    }

    prometheus::Gauge&
    executingGauge() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::internal_cgo_executing_task_total_search
                   : milvus::monitor::internal_cgo_executing_task_total_load;
    }

    prometheus::Histogram&
    queueDurationHist() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::internal_cgo_queue_duration_seconds_search
                   : milvus::monitor::internal_cgo_queue_duration_seconds_load;
    }

    prometheus::Histogram&
    executeDurationHist() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::
                         internal_cgo_execute_duration_seconds_search
                   : milvus::monitor::
                         internal_cgo_execute_duration_seconds_load;
    }

    prometheus::Counter&
    cancelBeforeCounter() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::
                         internal_cgo_cancel_before_execute_total_search
                   : milvus::monitor::
                         internal_cgo_cancel_before_execute_total_load;
    }

    prometheus::Counter&
    cancelDuringCounter() const {
        return pool_type_ == PoolType::kSearch
                   ? milvus::monitor::
                         internal_cgo_cancel_during_execute_total_search
                   : milvus::monitor::
                         internal_cgo_cancel_during_execute_total_load;
    }

    const PoolType pool_type_;
    std::chrono::steady_clock::time_point time_point_;
    Duration queue_duration_;
    Duration execute_duration_;
    bool cancelled_before_execute_{false};
    bool cancelled_during_execute_{false};
};

// FutureResult is a struct that represents the result of the future.
class FutureResult {
 public:
    explicit FutureResult(PoolType pool_type = PoolType::kSearch)
        : result(nullptr), status(), metrics(pool_type) {
    }
    void* result;
    CStatus status;
    Metrics<std::chrono::microseconds> metrics;
};

/// @brief a virtual class that represents a future can be polymorphic called by CGO code.
/// implemented by Future template.
class IFuture {
 public:
    /// @brief cancel the future with the given exception.
    /// After cancelled is called, the underlying async function will receive cancellation.
    /// It just a signal notification, the cancellation is handled by user-defined.
    /// If the underlying async function ignore the cancellation signal, the Future is still blocked.
    virtual void
    cancel() = 0;

    /// @brief check if the future is ready or canceled.
    /// @return true if the future is ready or canceled, otherwise false.
    virtual bool
    isReady() = 0;

    /// @brief register a callback that will be called when the future is ready or future has been ready.
    virtual void
    registerReadyCallback(CUnlockGoMutexFn unlockFn, CLockedGoMutex* mutex) = 0;

    /// @brief get the result of the future. it must be called if future is ready.
    /// the first element of the pair is the result,
    /// the second element of the pair is the exception.
    /// !!! It can only be called once,
    /// and the result need to be manually released by caller after these call.
    virtual std::pair<void*, CStatus>
    leakyGet() = 0;

    /// @brief leaked future object created by method `Future<R>::createLeakedFuture` can be droped by these method.
    static void
    releaseLeakedFuture(IFuture* future) {
        delete future;
    }

    virtual ~IFuture() = default;
};

/// @brief a helper function to throw a FutureCancellation exception if the token is cancelled.
static inline void
throwIfCancelled(const folly::CancellationToken& token) {
    if (token.isCancellationRequested()) {
        throw folly::FutureCancellation();
    }
}

/// @brief Future is a class that bound a future with a result for
/// using by cgo.
/// @tparam R is the return type of the producer function.
template <class R>
class Future : public IFuture {
 public:
    /// @brief do a async operation which will produce a result.
    /// fn returns pointer to R (leaked, default memory allocator) if it is success, otherwise it will throw a exception.
    /// returned result or exception will be handled by consumer side.
    template <typename Fn,
              typename = std::enable_if<
                  std::is_invocable_r_v<R*, Fn, folly::CancellationToken>>>
    static std::unique_ptr<Future<R>>
    async(folly::Executor::KeepAlive<> executor,
          int priority,
          Fn&& fn,
          PoolType pool_type = PoolType::kSearch) noexcept {
        auto future = std::make_unique<Future<R>>(pool_type);
        // setup the interrupt handler for the promise.
        future->setInterruptHandler();
        // start async function.
        future->asyncProduce(executor, priority, std::forward<Fn>(fn));
        // register consume callback function.
        future->registerConsumeCallback(executor, priority);
        return future;
    }

    /// use `async`.
    explicit Future(PoolType pool_type = PoolType::kSearch)
        : metrics_(pool_type),
          ready_(std::make_shared<Ready<LeakyResult<R>>>()),
          promise_(std::make_shared<folly::SharedPromise<R*>>()),
          cancellation_source_() {
    }

    Future(const Future<R>&) = delete;

    Future(Future<R>&&) noexcept = default;

    Future&
    operator=(const Future<R>&) = delete;

    Future&
    operator=(Future<R>&&) noexcept = default;

    /// @brief see `IFuture::cancel`
    void
    cancel() noexcept override {
        promise_->getSemiFuture().cancel();
    }

    /// @brief see `IFuture::registerReadyCallback`
    void
    registerReadyCallback(CUnlockGoMutexFn unlockFn,
                          CLockedGoMutex* mutex) noexcept override {
        ready_->callOrRegisterCallback(
            [unlockFn = unlockFn, mutex = mutex]() { unlockFn(mutex); });
    }

    /// @brief see `IFuture::isReady`
    bool
    isReady() noexcept override {
        return ready_->isReady();
    }

    /// @brief see `IFuture::leakyGet`
    std::pair<void*, CStatus>
    leakyGet() noexcept override {
        auto result = std::move(*ready_).getValue();
        return result.leakyGet();
    }

 private:
    /// @brief set the interrupt handler for the promise used in async produce arm.
    void
    setInterruptHandler() {
        promise_->setInterruptHandler([cancellation_source =
                                           cancellation_source_,
                                       ready = ready_](
                                          const folly::exception_wrapper& ew) {
            // 1. set the result to perform a fast fail.
            // 2. set the cancellation to the source to notify cancellation to the consumers.
            ew.handle(
                [&](const folly::FutureCancellation& e) {
                    cancellation_source.requestCancellation();
                },
                [&](const folly::FutureTimeout& e) {
                    cancellation_source.requestCancellation();
                });
        });
    }

    /// @brief do the R produce operation in async way.
    template <typename Fn,
              typename... Args,
              typename = std::enable_if<
                  std::is_invocable_r_v<R*, Fn, folly::CancellationToken>>>
    void
    asyncProduce(folly::Executor::KeepAlive<> executor, int priority, Fn&& fn) {
        // start produce process async.
        auto cancellation_token = cancellation_source_.getToken();
        auto runner = [fn = std::forward<Fn>(fn),
                       cancellation_token = std::move(cancellation_token),
                       this]() {
            // pre check the cancellation token
            if (cancellation_token.isCancellationRequested()) {
                metrics_.withEarlyCancel();
                throw folly::FutureCancellation();
            }

            // start the execution guard.
            Metrics<std::chrono::microseconds>::ExecutionGuard executionGuard(
                metrics_);

            try {
                return fn(cancellation_token);
            } catch (const folly::FutureCancellation& e) {
                metrics_.withDuringCancel();
                throw e;
            } catch (const milvus::SegcoreError& e) {
                // A cancellation raised inside execution does not always reach
                // here as a raw folly::FutureCancellation: the exec driver's
                // CALL_OPERATOR (and the load path) wrap it into a SegcoreError
                // carrying ErrorCode::FollyCancel so the sync and async CGO
                // boundaries agree on the code. Recognize that wrapped form so
                // a mid-execution cancellation is still counted as a during-
                // execute cancel instead of being dropped.
                if (e.get_error_code() == milvus::ErrorCode::FollyCancel) {
                    metrics_.withDuringCancel();
                }
                throw;  // rethrow the exception to the consumer side.
            } catch (...) {
                throw;  // rethrow the exception to the consumer side.
            }
        };

        // the runner is executed may be executed in different thread.
        // so manage the promise with shared_ptr.
        auto thenRunner = [promise = promise_, runner = std::move(runner)](
                              auto&&) { promise->setWith(std::move(runner)); };
        folly::makeSemiFuture().via(executor, priority).then(thenRunner);
    }

    /// @brief async consume the result of the future.
    void
    registerConsumeCallback(folly::Executor::KeepAlive<> executor,
                            int priority) noexcept {
        // set up the result consume arm and exception consume arm.
        promise_->getSemiFuture()
            .via(executor, priority)
            .thenValue(
                [ready = ready_](R* r) { ready->setValue(LeakyResult<R>(r)); })
            .thenError(folly::tag_t<folly::FutureCancellation>{},
                       [ready = ready_](const folly::FutureCancellation& e) {
                           ready->setValue(
                               LeakyResult<R>(milvus::FollyCancel, e.what()));
                       })
            .thenError(folly::tag_t<folly::FutureException>{},
                       [ready = ready_](const folly::FutureException& e) {
                           ready->setValue(LeakyResult<R>(
                               milvus::FollyOtherException, e.what()));
                       })
            .thenError(folly::tag_t<milvus::SegcoreError>{},
                       [ready = ready_](const milvus::SegcoreError& e) {
                           ready->setValue(LeakyResult<R>(
                               static_cast<int>(e.get_error_code()), e.what()));
                       })
            .thenError(
                folly::tag_t<std::exception>{},
                [ready = ready_](const std::exception& e) {
                    ready->setValue(LeakyResult<R>(
                        milvus::UnexpectedError,
                        fmt::format("{} :{}", typeid(e).name(), e.what())));
                });
    }

 private:
    Metrics<std::chrono::microseconds> metrics_;
    std::shared_ptr<Ready<LeakyResult<R>>> ready_;
    std::shared_ptr<folly::SharedPromise<R*>> promise_;
    folly::CancellationSource cancellation_source_;
};
};  // namespace milvus::futures