#include "async_coro.h"

#include <balancer/kernel/dns/common.h>
#include <balancer/kernel/dns/async_coro_requests.h>
#include <balancer/kernel/dns/async_wrapper.h>

#include <util/generic/hash_set.h>


namespace NSrvKernel::NDns {

namespace {
const auto DefaultCoroSleep = TDuration::MilliSeconds(10);
}

// ----------------------------------------------------------------------------
// TPollEvent
//
TPollEvent::TPollEvent(TContExecutor& e, SOCKET s, int what, PollCallback callback)
    : IPollEvent(s, what)
    , Executor_(e)
    , Callback_(callback)
{
    Executor_.Poller()->Schedule(this);
}

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

void TPollEvent::OnPollEvent(int) noexcept
{
    Callback_(Fd());
}


// ----------------------------------------------------------------------------
// TAsyncCoro
//
TAsyncCoro::TAsyncCoro(TContExecutor& e, const TOptions& opts, ICounters& counters, const TDnsServer& server)
    : Executor_(e)
    , Pool_(TDefaultAllocator::Instance())
    , Requests_(MakeHolder<TCoroRequests>(e))
    , Counters_(counters)
{
    auto socketStateCb = [this](SOCKET s, bool read, bool write) noexcept { OnStateChange(s, read, write); };
    AsyncDns_ = MakeHolder<TAsyncWrapper>(socketStateCb, opts, server);
}

TAsyncCoro::~TAsyncCoro()
{
    if (CoroAsyncDns_) {
        CoroAsyncDns_->Cancel();
        AsyncDns_->CancelRequests(); // sync
        Executor_.Running()->Join(CoroAsyncDns_, TInstant::Now());
    }
    Y_VERIFY(!CoroAsyncDns_, "coroutine pointer should be reseted on destroy");
}

bool TAsyncCoro::Resolve(const TString& host, int family, TInstant deadline, IEntry& result) noexcept
{
    Y_ASSERT(!host.empty());
    Y_ASSERT(family == AF_UNSPEC || family == AF_INET6 || family == AF_INET);
    auto scheduler = [this](IRequest& req) noexcept {
        Y_VERIFY(AsyncDns_);
        AsyncDns_->AsyncResolve(req.Host().c_str(), req.Family(), &req);
        ResolveCont()->ReSchedule();
    };
    Requests_->ExecuteRequest(host, family, deadline, scheduler, Counters_, result);
    return result.HasValidResult();
}

TCont* TAsyncCoro::ResolveCont() noexcept
{
    if (!CoroAsyncDns_) {
        CoroAsyncDns_ = Executor_.Create<TAsyncCoro, &TAsyncCoro::RunAsyncDns>(this, "async_dns_resolve");
    }

    return CoroAsyncDns_;
}

// TODO (alexeylaptev) : remove when BALANCER-3093 is fixed
void TAsyncCoro::StubForHttp2(TCont*) {
}

void TAsyncCoro::RunAsyncDns(TCont*) noexcept
{
    while (true) {
        if (!Requests_->Size()) {
            Executor_.Create<TAsyncCoro, &TAsyncCoro::StubForHttp2>(this, "http2_stub"); // see BALANCER-3093
            break;
        }

        if (!CoroAsyncDns_ || CoroAsyncDns_->Cancelled()) {
            Y_ASSERT(CoroAsyncDns_);
            Executor_.Create<TAsyncCoro, &TAsyncCoro::StubForHttp2>(this, "http2_stub2"); // see BALANCER-3093
            break;
        }

        Y_VERIFY(!Events_.Empty(), "events are not empty");

        // TODO (alexeylaptev) : move timeout to config
        if (CoroAsyncDns_ && !CoroAsyncDns_->Cancelled()
            && Executor_.Running()->SleepT(DefaultCoroSleep) == ETIMEDOUT)
        {
            AsyncDns_->ProcessNone();
        }

        if (CoroAsyncDns_ && !CoroAsyncDns_->Cancelled() && Requests_->Size()) {
            CoroAsyncDns_->Yield();
        }
    }

    CoroAsyncDns_ = nullptr;
}

void TAsyncCoro::OnStateChange(SOCKET s, bool read, bool write) noexcept try
{
    static const ui16 WHAT[] = {
            0, CONT_POLL_READ, CONT_POLL_WRITE, CONT_POLL_READ | CONT_POLL_WRITE};

    const int what = WHAT[((size_t)read) | (((size_t)write) << 1)];

    EventHolder& ev = Sockets_.Get(s);

    if (ev && !what) {
        ev.Destroy();
        return;
    }

    auto cb = [this](SOCKET s) {
        if (AsyncDns_) {
            AsyncDns_->ProcessSocket(s);
        }
    };

    if (!ev) {
        ev.Reset(new(&Pool_) TPollEvent(Executor_, s, what, cb));
        Events_.PushBack(ev.Get());
    } else {
        if (what) {
            if (ev->What() != what) {
                ev.Reset(new(&Pool_) TPollEvent(Executor_, s, what, cb));
                Events_.PushBack(ev.Get());
            }
        }
    }

    if (CoroAsyncDns_) {
        CoroAsyncDns_->ReSchedule();
    }
} catch(...) {
}

THolder<IAsyncCoro> MakeContAsyncResolver(TContExecutor* e, const TOptions& opts, ICounters& counters, const TDnsServer& server)
{
    return MakeHolder<TAsyncCoro>(*e, opts, counters, server);
}

} // namespace NSrvKernel::NDns