#pragma once

#include "resolver_face.h"

#include <balancer/kernel/thread/threadedqueue.h>
#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/helpers/exceptionless.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/stats/manager.h>
#include <library/cpp/coroutine/engine/events.h>

#include <util/digest/city.h>
#include <util/generic/noncopyable.h>
#include <util/generic/ptr.h>
#include <util/generic/scope.h>
#include <util/network/ip.h>

#include <variant>

class TContExecutor;
class IThreadFactory;

namespace NSrvKernel::NDns {
    class TThreadedResolver;

    template <class T>
    struct TAddressHash {
        T& operator[](const TResolveInfo& info) noexcept {
            ui64 val = (CityHash64(info.Host) << 16) | ((ui64)info.Port);
            if (info.Type == TResolveInfo::NUMERIC) {
                return Numeric_[val];
            } else {
                return General_[val];
            }
        }

        void Reset() noexcept {
            General_.clear();
            Numeric_.clear();
        }

    private:
        THashMap<ui64, T> General_;
        THashMap<ui64, T> Numeric_;
    };

    class TContResolver : public IResolver, public TNonCopyable {
    private:
        class TAddressFuture {
        public:
            TAddressFuture(TContExecutor* executor, TThreadedResolver* slave,
                           const TResolveInfo& info, TDuration recordTtl)
                : WaitEvent_(executor)
                , Executor_(executor)
                , Slave_(slave)
                , Info_(info)
                , RecordTtl_(recordTtl)
                , ResultingAddress_(Y_MAKE_ERROR(yexception{} << "TContExecutor caching error"))
            {
                Y_ASSERT(Slave_);
                StartCont();
            }
            ~TAddressFuture();

            TErrorOr<TSockAddrInfo> GetAddress(TInstant deadline, TStatsCounters& counters) noexcept;

            void Cancel() {
                Task_ = {};
            }

        private:
            void StartCont();
            bool DeadlinePassed() const noexcept {
                return RecordTtl_ != TDuration::Max() && Now() > TtlDeadline_;
            }

        private:
            TContSimpleEvent WaitEvent_;
            TContExecutor* Executor_ = nullptr;
            TThreadedResolver* Slave_ = nullptr; // not owned
            TResolveInfo Info_;
            TDuration RecordTtl_ = TDuration::Max();
            TInstant TtlDeadline_ = TInstant::Max();
            TErrorOr<TSockAddrInfo> ResultingAddress_;
            TCoroutine Task_;
        };

        using TAddressFuturePtr = THolder<TAddressFuture>;

    public:
        TContResolver(TContExecutor* executor, TStatsCounters& counters, TThreadedQueue* queue, TDuration timeout, TDuration recordTtl);
        ~TContResolver() override;

    private:
        TErrorOr<TSockAddrInfo> DoResolve(const TResolveInfo& info, bool, TInstant deadline) override {
            const auto start = TInstant::Now();
            TAddressFuturePtr &future = Hash_[info];

            if (!future) {
                future = MakeHolder<TAddressFuture>(Executor_, Slave_.get(), info, RecordTtl_);
            }

            Y_DEFER {
                const auto resolveTime = TInstant::Now() - start;
                Counters_.DnsResolveTime.AddValue(resolveTime.MicroSeconds());
            };
            return future->GetAddress(deadline, Counters_);
        }

        void DoResetCache() override {
            Counters_.DnsResetCache.Add(1);
            Hash_.Reset();
        }

    private:
        TContExecutor* Executor_ = nullptr;
        TStatsCounters& Counters_;
        std::unique_ptr<TThreadedResolver> Slave_;
        TAddressHash<TAddressFuturePtr> Hash_;
        TDuration RecordTtl_ = TDuration::Max();
    };

}
