#pragma once

#include <library/cpp/coroutine/engine/impl.h>

#include <util/generic/scope.h>

namespace NSrvKernel {
    class TPollerEvent:
        public NCoro::IPollEvent
    {
    public:
        TPollerEvent(TContExecutor& executor, SOCKET fd, ui16 what) noexcept
            : IPollEvent(fd, what)
            , Executor_(executor)
        {
            Executor_.Poller()->Schedule(this);
        }

        ~TPollerEvent() {
            Executor_.Poller()->Remove(this);
        }

        int Wait(TCont* cont, TInstant deadline) {
            Y_VERIFY(cont->Executor() == &Executor_);

            if (cont->Cancelled()) {
                return ECANCELED;
            }

            NCoro::TContPollEvent contEvent(cont, deadline);
            {
                ContEvent_ = &contEvent;
                Y_DEFER {
                    ContEvent_ = nullptr;
                };

                Executor_.RegisterInWaitQueue(ContEvent_);
                cont->Switch();
            }

            if (cont->Cancelled()) {
                return ECANCELED;
            }

            return contEvent.Status();
        }

        void SetUnwaitedEventCallback(const std::function<void(int)>& cb) {
            UnwaitedEventCb_ = cb;
        }

        void OnPollEvent(int status) noexcept override {
            if (ContEvent_) {
                ContEvent_->Wake(status);
            } else {
                UnwaitedEventCb_(status);
            }
        }
    private:
        TContExecutor& Executor_;
        NCoro::TContPollEvent* ContEvent_ = nullptr;
        std::function<void(int)> UnwaitedEventCb_;
    };
}
