#include "async_cache.h"

#include "async_coro.h"
#include "async_wrapper.h"
#include "resolver_face.h"

#include <balancer/kernel/dns/common.h>
#include <balancer/kernel/helpers/exceptionless.h>
#include <balancer/kernel/net/addr.h>

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


namespace NSrvKernel::NDns {

TAsyncCache::TAsyncCache(TContExecutor* e, const TResolverConfig& config,
                         THolder<IAsyncCoro>&& resolver, THolder<IEntryBuilder>&& builder,
                         ICounters& counters)
    : Executor_(e)
    , Config_(config)
    , Counters_(counters)
    , Resolver_(std::move(resolver))
    , EntryBuilder_(std::move(builder))
{
    Y_VERIFY(e);
    Y_VERIFY(Resolver_ && EntryBuilder_);
}

TErrorOr<TSockAddrInfo> TAsyncCache::LookupOrResolve(const TResolveInfo& info, bool onlyCache, TInstant deadline) {
    TEntryHolder& entry = Storage_[info.Host];
    if (!entry) {
        entry = EntryBuilder_->GetNewEntry(Executor_, Config_.RecordTtl);
    }

    const TInstant resolveTimeout = onlyCache ? TInstant::Now() : deadline;

    bool run = true;
    while (run) {
        Y_ASSERT(entry);
        switch (entry->GetState()) {
            case RS_NOT_RESOLVED:
                entry->SetInProgress();
                if (!Resolve(info.Host, info.Family, resolveTimeout, *entry)) {
                    return Y_MAKE_ERROR(TNetworkResolutionError{EAI_FAIL} << (onlyCache ? " not in cache" : " resolve failed"));
                }
                run = false; // states were set inside Resolve(), just break the loop here
                break;
            case RS_IN_PROGRESS: {
                if (onlyCache) {
                    Counters_.AddCacheError();
                    return Y_MAKE_ERROR(TNetworkResolutionError{EAI_FAIL} << " not in cache");
                }
                int waitResultCode = 0;
                if (!entry->WaitD(deadline, waitResultCode)) {
                    Counters_.AddCacheError();
                    if (waitResultCode != ECANCELED) {
                        return Y_MAKE_ERROR(TNetworkResolutionError{EAI_FAIL} << " wait failed with code " << ::ToString(waitResultCode));
                    } else {
                        return Y_MAKE_ERROR(yexception{} << " wait was cancelled");
                    }
                }
                continue;
            }
            case RS_RESOLVED:
                if (entry->IsExpired(Config_.RecordTtl) && !entry->HasRequestsInProgress()) {
                    if (!Resolve(info.Host, info.Family, resolveTimeout, *entry)) {
                        Counters_.AddCacheResolve(); // use expired from cache
                    }
                } else {
                    Counters_.AddCacheResolve();
                }
                run = false;
                break;
            case RS_FAILED:
                if (entry->ResetFailedIfExpired(Config_.ErrorTtl)) {
                    continue;
                }
                Counters_.AddCacheError();
                return Y_MAKE_ERROR(TNetworkResolutionError{EAI_FAIL} << " resolve failed");
            default:
                Y_ASSERT(false);
                Counters_.AddResolveError();
                return Y_MAKE_ERROR(yexception{} << "Bad state in DNS cache, shouldn't not get here");
        };
    }

    Y_ASSERT(entry->GetState() == RS_RESOLVED);

    if (!entry->HasValidResult()) {
        Y_ASSERT(false); // shouldn't get resolved status if no addresses were provided
        Counters_.AddResolveError();
        return Y_MAKE_ERROR(yexception{} << "Failed to resolve - invalid addresses");
    }

    TSockAddrInfo returnResult = entry->GetAddrInfo();
    SetPort(returnResult, info.Port);
    return std::move(returnResult);
}

bool TAsyncCache::Resolve(const TString& host, int family, TInstant deadline, IEntry& result) noexcept
{
    return Resolver_->Resolve(host, family, deadline, result);
}

void TAsyncCache::SetPort(TSockAddrInfo& result, uint16_t port) const noexcept
{
    for (auto& i : result.Addresses) {
        i.SetPort(port);
    }
}

}