#include <balancer/kernel/http2/client/pinger.h>

#include <library/cpp/testing/gtest/gtest.h>

using namespace NBalancerClient;
using namespace NSrvKernel;
using namespace testing;

class TTime: public NCoro::ITime {
  public:

    void AddNow(TDuration diff) {
        Now_ += diff;
    }

    TInstant Now() override {
        return Now_;
    }

  private:
    TInstant Now_ = TInstant::Now();
};

MATCHER_P(IsBackendError, matcher, TStringBuilder{} << DescribeMatcher<NSrvKernel::TBackendError>(matcher, negation) <<  " and should be of type " << TypeName<NSrvKernel::TBackendError>()) {
    *result_listener << "has type " << TypeName(arg) << ", ";
    if (auto backendError = arg.template GetAs<NSrvKernel::TBackendError>()) {
        return ExplainMatchResult(matcher, *backendError, result_listener);
    }
    return false;
}

MATCHER_P(IsBackendErrorWithMsg, value, TStringBuilder{} << "should contain string '" << value << "'" ) {
    *result_listener << "has message \"" << arg.ErrorMessage() << "\"";
    return arg.ErrorMessage().Contains(value);;
}

class TPingTestsEnv {
  public:
    StrictMock<MockFunction<TError(TInstant)>> PingMock;
    StrictMock<MockFunction<void (TError)>> CancelMock;
    NiceMock<MockFunction<void (TPingContext::EState)>> StateMock;
    TTime Time;
    TContExecutor Executor;
    TPinger::TOptions Options;
    TPinger Pinger;
    TPingContext::TOptions ContextOptions;
    TPingContext PingContext;

    TPingTestsEnv(size_t maxSuccessivePings = 0)
        : Executor{1024*1024, IPollerFace::Default(), nullptr, nullptr, NCoro::NStack::EGuard::Canary, Nothing(), &Time}
        , Options{
            .ContExecutor = &Executor,
            .TimeWithoutDataBeforePing = TDuration::Seconds(1),
        }
        , Pinger{Options}
        , ContextOptions{
            .Pinger = Pinger,
            .ContExecutor = &Executor,
            .PingHandler = PingMock.AsStdFunction(),
            .CancelHandler = CancelMock.AsStdFunction(),
            .StateChangeHandler = StateMock.AsStdFunction(),
            .PingTimeout = TDuration::Seconds(1),
            .MaxSuccessivePings = maxSuccessivePings
        }
        , PingContext{ContextOptions}
    {
        Pinger.Shedule(PingContext);
    }

    void Cancel() {
        Pinger.Cancel();
        PingContext.Cancel();
        Executor.Execute();
    }
};

TEST(Pings, NotStarted) {
    TPinger pinger{{}};
    TPingContext pingContext{TPingContext::TOptions{
        .Pinger = pinger
    }};
    ASSERT_EQ(pingContext.GetState(), TPingContext::EState::NotStarted);
}

TEST(Pings, IdleTimer) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.CancelMock, Call).Times(0);
    EXPECT_CALL(env.PingMock, Call).Times(0);
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing/2);
        RunningCont()->Yield();
        env.Executor.Pause();
    });
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::IdleTimer);

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Cancel();
}

TEST(Pings, PingTimeout) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout)))
        .WillOnce(Return(ByMove(TError{})));
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.Time.AddNow(env.ContextOptions.PingTimeout);
    });
    EXPECT_CALL(env.CancelMock, Call(IsBackendError(IsBackendErrorWithMsg("ping failed")))).WillOnce(Return());
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingFailed)).WillOnce(Return());
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::PingFailed);
}

TEST(Pings, PingSendFailure) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout)))
        .WillOnce(Return(ByMove(Y_MAKE_ERROR(yexception{} << "some error"))));
    EXPECT_CALL(env.CancelMock, Call(IsBackendError(IsBackendErrorWithMsg("ping send failed")))).WillOnce(Return());
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingFailed));
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::PingFailed);
}

TEST(Pings, LongSendPing) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout))).WillOnce([&](TInstant) {
        env.Time.AddNow(env.ContextOptions.PingTimeout*2);
        return TError{};
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer));
    EXPECT_CALL(env.CancelMock, Call(IsBackendError(IsBackendErrorWithMsg("ping failed"))))
        .WillOnce(Return());
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingFailed));
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::PingFailed);
}

TEST(Pings, SuccessPing) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout)))
        .WillOnce(Return(ByMove(TError{})));
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.PingContext.PingAck();
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Executor.Pause();
    });
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetSuccessivePings(), size_t{1});
    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::IdleTimer);

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Cancel();
}

TEST(Pings, SuccessivePings) {
    TPingTestsEnv env{1};
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout)))
        .WillOnce(Return(ByMove(TError{})));
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.PingContext.PingAck();
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::LimitReached)).WillOnce([&]() {
        env.Executor.Pause();
    });
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetSuccessivePings(), size_t(1));
    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::LimitReached);

    env.PingContext.ResetSuccessivePings();

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Executor.Pause();
    });
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetSuccessivePings(), size_t(0));
    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::IdleTimer);

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Cancel();
}

TEST(Pings, CancelWaitingPingsReset) {
    TPingTestsEnv env{1};
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout)))
        .WillOnce(Return(ByMove(TError{})));
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.PingContext.PingAck();
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::LimitReached)).WillOnce([&]() {
        env.PingContext.Cancel();
    });
    EXPECT_CALL(env.CancelMock, Call).Times(0);
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Executor.Execute();
}

TEST(Pings, CancelWaitingPingTimer) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout))).WillOnce(Return(ByMove(TError{})));
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.PingContext.Cancel();
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Executor.Execute();
}

TEST(Pings, UnexpectedPing) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);
    });
    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.Options.TimeWithoutDataBeforePing + env.ContextOptions.PingTimeout))).WillOnce([&]() {
        env.PingContext.PingAck();
        return TError{};
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::PingTimer)).WillOnce([&]() {
        env.Executor.Pause();
    });
    env.Executor.Execute();

    ASSERT_EQ(env.PingContext.GetSuccessivePings(), size_t(0));
    ASSERT_EQ(env.PingContext.GetState(), TPingContext::EState::PingTimer);

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Cancel();
}

TEST(Pings, PostponePing) {
    TPingTestsEnv env;
    InSequence sequence;

    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::IdleTimer)).WillOnce([&]() {
        env.Executor.Pause();
    });
    env.Executor.Execute();

    env.Time.AddNow(env.Options.TimeWithoutDataBeforePing/2);
    env.PingContext.DataReceived();
    env.Time.AddNow(env.Options.TimeWithoutDataBeforePing);

    EXPECT_CALL(env.PingMock, Call(Eq(env.Time.Now() + env.ContextOptions.PingTimeout))).WillOnce([&](){
        env.PingContext.Cancel();
        return TError{};
    });
    EXPECT_CALL(env.StateMock, Call(TPingContext::EState::Cancelled)).WillOnce(Return());
    env.Executor.Execute();
}
