#include <balancer/kernel/dns/async_entry.h>

#include <library/cpp/testing/gmock_in_unittest/gmock.h>
#include <library/cpp/testing/unittest/registar.h>

#include <thread>

using namespace ::testing;

namespace NSrvKernel::NDns::NTests {

namespace {
constexpr TDuration DefaultRecordTtl = TDuration::MilliSeconds(1000 * 5);
}

class TEntryFixture : public NUnitTest::TBaseFixture {
public:
    TEntryFixture()
        : Executor_(64 * 1024)
        , Addr_(TIpAddr::FromIp("::1").GetOrThrow(), 0)
        , Addr2_(TIpAddr::FromIp("127.0.0.1").GetOrThrow(), 0)
        , Addr3_(TIpAddr::FromIp("8.8.8.8").GetOrThrow(), 0)
    {
        RecreateEntry(DefaultRecordTtl);
    }

protected:
    void RecreateEntry(TDuration recordTtl)
    {
        Entry_ = MakeHolder<TEntry>(&Executor_, recordTtl);
    }

protected:
    TContExecutor Executor_;
    THolder<TEntry> Entry_;
    TSockAddr Addr_;
    TSockAddr Addr2_;
    TSockAddr Addr3_;
};

Y_UNIT_TEST_SUITE_F(DnsAsyncEntry, TEntryFixture) {

Y_UNIT_TEST_F(DefaultCreated_NotResolvedState, TEntryFixture)
{
    EXPECT_EQ(Entry_->GetState(), RS_NOT_RESOLVED);
}

Y_UNIT_TEST_F(AllRequestsFailed_FailedState, TEntryFixture)
{
    EXPECT_EQ(Entry_->GetState(), RS_NOT_RESOLVED);

    Entry_->SetInProgress();
    EXPECT_EQ(Entry_->GetState(), RS_IN_PROGRESS);

    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    EXPECT_EQ(Entry_->GetState(), RS_IN_PROGRESS);

    Entry_->SetFailedIfDone();

    EXPECT_EQ(Entry_->GetState(), RS_FAILED);

    Entry_->SetFailedIfDone();
    EXPECT_EQ(Entry_->GetState(), RS_FAILED);
}

Y_UNIT_TEST_F(ResultAdded_ResolvedState, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->AddRequestInProgress();

    Entry_->AddResult(std::move(Addr_));
    EXPECT_EQ(Entry_->GetState(), RS_RESOLVED);

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    EXPECT_EQ(Entry_->GetState(), RS_RESOLVED);
}

Y_UNIT_TEST_F(ResetInProgressAsFailedState_InProgressState, TEntryFixture)
{
    Entry_->SetInProgress();

    EXPECT_FALSE(Entry_->ResetFailedIfExpired(TDuration::MilliSeconds(0)));
    EXPECT_EQ(Entry_->GetState(), RS_IN_PROGRESS);
}

Y_UNIT_TEST_F(ResetNotExpiredFailedState_FailedState, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->SetFailedIfDone();
    EXPECT_EQ(Entry_->GetState(), RS_FAILED);

    EXPECT_FALSE(Entry_->ResetFailedIfExpired(TDuration::MilliSeconds(1000 * 5)));
    EXPECT_EQ(Entry_->GetState(), RS_FAILED);
}

Y_UNIT_TEST_F(ResetExpiredFailedState_NotResolvedState, TEntryFixture)
{
    Entry_->SetInProgress();
    EXPECT_TRUE(Entry_->SetFailedIfDone());
    EXPECT_EQ(Entry_->GetState(), RS_FAILED);

    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    const auto ttl = TDuration::MilliSeconds(1);
    EXPECT_TRUE(Entry_->IsExpired(ttl));

    EXPECT_TRUE(Entry_->ResetFailedIfExpired(ttl));
    EXPECT_EQ(Entry_->GetState(), RS_NOT_RESOLVED);
}

Y_UNIT_TEST_F(SeveralResultsAdded_AllStored, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();
    Entry_->SetFailedIfDone();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();

    TVector<TSockAddr> origin{Addr_, Addr2_, Addr3_};
    Entry_->AddResult(std::move(Addr_));
    Entry_->AddResult(std::move(Addr2_));
    Entry_->AddResult(std::move(Addr3_));

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    TSockAddrInfo result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 3);

    for (size_t i : {0, 1, 2}) {
        EXPECT_EQ(result.Addresses[i], origin[i]);
    }
}

Y_UNIT_TEST_F(SeveralResultsAdded_DuplicatesRemovedAllStored, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();
    Entry_->SetFailedIfDone();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();

    TVector<TSockAddr> origin{Addr_, Addr2_};
    Entry_->AddResult(std::move(Addr_));
    Entry_->AddResult(std::move(Addr2_));
    Entry_->AddResult(std::move(Addr_));
    Entry_->AddResult(std::move(Addr2_));
    Entry_->AddResult(std::move(Addr2_));

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    TSockAddrInfo result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 2);

    for (size_t i : {0, 1}) {
        EXPECT_EQ(result.Addresses[i], origin[i]);
    }
}

Y_UNIT_TEST_F(AddNewDataAfterRecordTtl_OldDataRemoved, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();

    TVector<TSockAddr> origin{Addr_, Addr2_};
    Entry_->AddResult(std::move(Addr_));
    Entry_->AddResult(std::move(Addr2_));

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    TSockAddrInfo result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 2);
    for (size_t i : {0, 1}) {
        EXPECT_EQ(result.Addresses[i], origin[i]);
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(DefaultRecordTtl.MilliSeconds()));

    Entry_->AddRequestInProgress();

    TVector<TSockAddr> originNew{Addr3_};
    Entry_->AddResult(std::move(Addr3_));

    Entry_->SetFailedIfDone();

    result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 1);
    EXPECT_EQ(result.Addresses[0], originNew[0]);
}

Y_UNIT_TEST_F(EmptyDataAddedAfterRecordTtl_OldDataNotRemoved, TEntryFixture)
{
    Entry_->SetInProgress();
    Entry_->AddRequestInProgress();
    Entry_->AddRequestInProgress();

    TVector<TSockAddr> origin{Addr_, Addr2_};
    Entry_->AddResult(std::move(Addr_));
    Entry_->AddResult(std::move(Addr2_));

    Entry_->SetFailedIfDone();
    Entry_->SetFailedIfDone();

    TSockAddrInfo result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 2);
    for (size_t i : {0, 1}) {
        EXPECT_EQ(result.Addresses[i], origin[i]);
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(DefaultRecordTtl.MilliSeconds()));

    Entry_->AddRequestInProgress();

    Entry_->AddResult(TSockAddr{});

    Entry_->SetFailedIfDone();

    result = Entry_->GetAddrInfo();
    UNIT_ASSERT(result.Addresses.size() == 2);
    for (size_t i : {0, 1}) {
        EXPECT_EQ(result.Addresses[i], origin[i]);
    }
}

} // ens of fixture pack
}