#pragma once

#include "coroutine.h"

#include <variant>

namespace NSrvKernel {

namespace NPrivate {

struct TCoroFutureAccessor;

}  // namespace NPrivate

/**
 * TCoroFuture class provides a mechanism to access the result of asynchronous operation.
 *
 * TCoroFuture is a result CoroAsync call.
 **/
template <class T>
class TCoroFuture {
public:
    TCoroFuture() = default;

    TCoroFuture(TCoroFuture&& rhs) noexcept
        : Result_(std::move(rhs.Result_))
        , Hook_(std::move(rhs.Hook_))
        , Coroutine_(std::move(rhs.Coroutine_))
    {
        if (Valid()) {
            Hook_->ResultPtr = &Result_;
        }
    }

    TCoroFuture& operator=(TCoroFuture&& rhs) noexcept {
        auto tmp = std::move(rhs);
        Swap(tmp);
        return *this;
    }

    void Swap(TCoroFuture& rhs) noexcept {
        if (Valid()) {
            if (rhs.Valid()) {
                DoSwap(Result_, rhs.Result_);
                DoSwap(Hook_, rhs.Hook_);
                DoSwap(Coroutine_, rhs.Coroutine_);
                Hook_->ResultPtr = &Result_;
                rhs.Hook_->ResultPtr = &rhs.Result_;
            } else {
                rhs.Swap(*this);
            }
        } else if (rhs.Valid()) {
            Result_ = std::move(rhs.Result_);
            Hook_ = std::move(rhs.Hook_);
            Coroutine_ = std::move(rhs.Coroutine_);
            Hook_->ResultPtr = &Result_;
        }
    }

    /**
     * Waits until result becomes available and retrieves it.
     *
     * According to the standard regarding std::future:
     * The behavior is undefined if Valid() is false before the call to this function.
     *
     * Valid() is false after a call to this method.
     **/
    T Get() {
        Wait();
        Hook_.Reset();

        return std::visit([](auto val) -> T {
            if constexpr (std::is_same_v<decltype(val), std::exception_ptr>) {
                std::rethrow_exception(std::move(val));
            } else {
                return val;
            }
        }, std::move(Result_));
    }

    /**
     * Checks if the future refers to a shared state (coro context).
     **/
    bool Valid() const noexcept {
        return Hook_ != nullptr;
    }

    /**
     * Blocks until the result becomes available.
     *
     * According to the standard regarding std::future:
     * The behavior is undefined if Valid() is false before the call to this function.
     *
     * Notice that this method is not marked const as opposed to std::future's wait() method
     * since we don't want to store the result in shared state on the heap.
     *
     * Valid() is true after a call to this method.
     **/
    void Wait() {
        Coroutine_.Join();
    }

    /**
     * Part of coroutine interface. Marks running coroutine as canceled.
     **/
    void Cancel() {
        Coroutine_.Cancel();
    }

private:
    friend struct NPrivate::TCoroFutureAccessor;

    std::variant<std::exception_ptr, T> Result_;
    struct THook {
        std::variant<std::exception_ptr, T>* ResultPtr = nullptr;
    };
    THolder<THook> Hook_;
    TCoroutine Coroutine_;
};

namespace NPrivate {

struct TCoroFutureAccessor {
    template <class R, class F, class... Args>
    [[nodiscard]] static TCoroFuture<R> CoroAsync(const char* name, TContExecutor* const executor,
                                                  F&& f, Args&&... args) {
        Y_ASSERT(executor != nullptr);

        TCoroFuture<R> ret;
        ret.Hook_ = MakeHolder<typename TCoroFuture<R>::THook>();
        ret.Hook_->ResultPtr = &ret.Result_;
        ret.Coroutine_ = TCoroutine{name, executor, [hook = ret.Hook_.Get()](
            std::decay_t<F> f, std::decay_t<Args>... args) {
            try {
                hook->ResultPtr->template emplace<1>(
                    std::invoke(std::move(f), std::move(args)...));
            } catch (...) {
                *hook->ResultPtr = std::current_exception();
            }
        }, std::forward<F>(f), std::forward<Args>(args)...};

        return ret;
    }
};

}  // namespace NPrivate

/**
 * Launches operation on new coroutine and returns TCoroFuture object.
 *
 * TCoroFuture object will eventually hold the result of that function call
 * (analogue of std::async for coroutines).
 *
 * Notice that arguments will be decayed to values. If you want to send references,
 * use reference wrappers (std::ref, std::cref).
 **/
template <class F, class... Args>
[[nodiscard]] TCoroFuture<std::invoke_result_t<std::decay_t<F>, std::decay_t<Args>...>>
CoroAsync(const char* name, TContExecutor* const executor, F&& f, Args&&... args) {
    using TResult = std::invoke_result_t<std::decay_t<F>, std::decay_t<Args>...>;

    static_assert(!std::is_void_v<TResult>,
                  "CoroAsync can't launch non-returning functions. Use TCoroutine instead.");

    static_assert(!std::is_reference_v<TResult>,
                  "CoroAsync can't launch functions that return references yet.");

    return NPrivate::TCoroFutureAccessor::CoroAsync<TResult>(
        name, executor, std::forward<F>(f), std::forward<Args>(args)...);
}

}  // namespace NSrvKernel
