#include <library/cpp/testing/unittest/registar.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/requester/requester.h>
#include <balancer/kernel/testing/conn_descr.h>
#include <balancer/kernel/testing/module_mock.h>
#include <balancer/kernel/testing/process_mock.h>

using namespace NSrvKernel;

using ::testing::_;
using ::testing::Invoke;
using ::testing::Return;

Y_UNIT_TEST_SUITE(NTestRequester) {
    const TString REQUEST_BODY("REQUEST_BODY");
    const TString RESPONSE_BODY("RESPONSE_BODY");
    const int STATUS_CODE = 418;
    THTTP2StreamDetails HTTP2StreamDetails_;

    void CheckConnDescr(
        const TConnDescr& origDescr,
        const TConnDescr& descr,
        const TRequest& request
    ) {
        UNIT_ASSERT_UNEQUAL(descr.Input, origDescr.Input);
        UNIT_ASSERT_UNEQUAL(descr.Output, origDescr.Output);
        UNIT_ASSERT_EQUAL(descr.Properties->Parent.RemoteAddress, origDescr.Properties->Parent.RemoteAddress);
        UNIT_ASSERT_EQUAL(descr.Properties->Parent.LocalAddress, origDescr.Properties->Parent.LocalAddress);
        UNIT_ASSERT_EQUAL(descr.ExtraAccessLog, origDescr.ExtraAccessLog);
        UNIT_ASSERT_EQUAL(descr.ErrorLog, origDescr.ErrorLog);
        UNIT_ASSERT_EQUAL(descr.Hash, origDescr.Hash);
        UNIT_ASSERT_EQUAL(descr.Request, &request);
    }

    void CheckGet(const TConnDescr& descr) {
        UNIT_ASSERT_EQUAL(descr.Request->Props().HTTP2, nullptr);
        UNIT_ASSERT(!descr.Request->Props().ContentLength);
        UNIT_ASSERT_EQUAL(descr.Request->Props().ChunkedTransfer, false);
    }

    void CheckPost(const TConnDescr& descr, bool chunked) {
        UNIT_ASSERT_EQUAL(descr.Request->Props().HTTP2, nullptr);
        if (chunked) {
            UNIT_ASSERT(!descr.Request->Props().ContentLength);
            UNIT_ASSERT_EQUAL(descr.Request->Props().ChunkedTransfer, true);
        } else {
            UNIT_ASSERT(descr.Request->Props().ContentLength);
            UNIT_ASSERT_EQUAL(*descr.Request->Props().ContentLength, REQUEST_BODY.size());
            UNIT_ASSERT_EQUAL(descr.Request->Props().ChunkedTransfer, false);
        }

        TChunkList chunk;
        TryRethrowError(descr.Input->Recv(chunk, TInstant::Max()));
        UNIT_ASSERT_EQUAL(chunk, REQUEST_BODY);
    }

    void CheckResponse(const TResponse& response) {
        UNIT_ASSERT_EQUAL(response.ResponseLine().StatusCode, STATUS_CODE);
    }

    void CheckResponseBody(TChunkList& responseBody) {
        UNIT_ASSERT_EQUAL(responseBody, RESPONSE_BODY);
    }

    void SendResponse(const TConnDescr& descr, bool sendBody) {
        TResponseLine responseLine;
        responseLine.StatusCode = STATUS_CODE;
        TResponse response{std::move(responseLine)};

        TryRethrowError(descr.Output->SendHead(std::move(response), false, TInstant::Max()));

        if (sendBody) {
            TryRethrowError(descr.Output->Send(TChunkList(RESPONSE_BODY), TInstant::Max()));
        }
        TryRethrowError(descr.Output->SendEof(TInstant::Max()));
    }

    TRequest GetRequest(bool head = false) {
        TRequest request;
        request.Props().HTTP2 = &HTTP2StreamDetails_;
        request.Props().ChunkedTransfer = true;
        request.Props().ContentLength = 12345;

        if (head) {
            request.RequestLine().Method = EMethod::HEAD;
        } else {
            request.RequestLine().Method = EMethod::GET;
        }
        return request;
    }

    Y_UNIT_TEST(RequesterGet) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckGet(descr);

                SendResponse(descr, false);
                return TError();
            }));

        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;

        TryRethrowError(requester.Request(std::move(request), response));
        CheckResponse(response);
    }

    Y_UNIT_TEST(RequesterGetNoResponse) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckGet(descr);

                SendResponse(descr, false);
                return TError();
            }));

        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;

        TryRethrowError(requester.Request(std::move(request)));
    }

    Y_UNIT_TEST(RequesterPost) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckPost(descr, false);

                SendResponse(descr, false);
                return TError();
            }));

        TChunkList requestBody(REQUEST_BODY);
        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;

        TryRethrowError(requester.Request(
            std::move(request),
            std::move(requestBody),
            false,
            response
        ));
        CheckResponse(response);
    }

    Y_UNIT_TEST(RequesterPostNoResponse) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckPost(descr, false);

                SendResponse(descr, false);
                return TError();
            }));

        TChunkList requestBody(REQUEST_BODY);
        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;

        TryRethrowError(requester.Request(
            std::move(request),
            std::move(requestBody),
            false
        ));
    }

    Y_UNIT_TEST(RequesterPostChunked) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckPost(descr, true);

                SendResponse(descr, false);
                return TError();
            }));

        TChunkList requestBody(REQUEST_BODY);
        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;

        TryRethrowError(requester.Request(std::move(request), std::move(requestBody), true, response));
        CheckResponse(response);
    }

    Y_UNIT_TEST(RequesterPostFromStreamRespBody) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckPost(descr, true);

                SendResponse(descr, true);
                return TError();
            }));

        TChunksInput chunksIo{TChunkList{REQUEST_BODY}};
        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;
        TChunkList responseBody;

        TryRethrowError(requester.Request(std::move(request), chunksIo, response, responseBody, true));
        CheckResponse(response);
    }

    Y_UNIT_TEST(RequesterGetRespBody) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckGet(descr);

                SendResponse(descr, true);
                return TError();
            }));

        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;
        TChunkList responseBody;

        TryRethrowError(requester.Request(std::move(request), response, responseBody));
        CheckResponse(response);
        CheckResponseBody(responseBody);
    }

    Y_UNIT_TEST(RequesterPostRespBody) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);
                CheckPost(descr, false);

                SendResponse(descr, true);
                return TError();
            }));

        TChunkList requestBody(REQUEST_BODY);
        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;
        TChunkList responseBody;

        TryRethrowError(requester.Request(
            std::move(request),
            std::move(requestBody),
            false,
            response,
            responseBody
        ));
        CheckResponse(response);
        CheckResponseBody(responseBody);
    }

    Y_UNIT_TEST(RequesterHeadTest) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest(true);

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                CheckConnDescr(origDescr.ConnDescr(), descr, request);

                SendResponse(descr, false);
                return TError();
            }));

        TRequester requester(moduleMock, origDescr.ConnDescr());
        TResponse response;
        TChunkList responseBody;

        TryRethrowError(requester.Request(
            std::move(request),
            response,
            responseBody
        ));
        CheckResponse(response);
        UNIT_ASSERT(responseBody.Empty());
    }

    Y_UNIT_TEST(AsyncRequester) {
        TLog errorLog;
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        TRequest request = GetRequest();

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                UNIT_ASSERT_EQUAL(descr.ErrorLog, &errorLog);
                UNIT_ASSERT_EQUAL(descr.Request, &request);

                SendResponse(descr, false);
                return TError();
            }));

        TAsyncRequester requester(moduleMock, &errorLog, process);
        TResponse response;

        TryRethrowError(requester.Requester().Request(std::move(request), response));
        CheckResponse(response);
    }

    Y_UNIT_TEST(RequesterTcpRequest) {
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;
        NSrvKernel::NTesting::TTestConnDescr origDescr(process);

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                UNIT_ASSERT_UNEQUAL(descr.Input, origDescr.ConnDescr().Input);
                UNIT_ASSERT_UNEQUAL(descr.Output, origDescr.ConnDescr().Output);
                UNIT_ASSERT_EQUAL(
                    descr.Properties->Parent.RemoteAddress, origDescr.ConnDescr().Properties->Parent.RemoteAddress);
                UNIT_ASSERT_EQUAL(
                    descr.Properties->Parent.LocalAddress, origDescr.ConnDescr().Properties->Parent.LocalAddress);
                UNIT_ASSERT_EQUAL(descr.ExtraAccessLog, origDescr.ConnDescr().ExtraAccessLog);
                UNIT_ASSERT_EQUAL(descr.ErrorLog, origDescr.ConnDescr().ErrorLog);
                UNIT_ASSERT_EQUAL(descr.Hash, origDescr.ConnDescr().Hash);
                UNIT_ASSERT_EQUAL(descr.Request, nullptr);

                TChunkList chunk;
                TryRethrowError(descr.Input->Recv(chunk, TInstant::Max()));
                UNIT_ASSERT_EQUAL(chunk, REQUEST_BODY);

                TryRethrowError(descr.Output->Send(TChunkList(RESPONSE_BODY), TInstant::Max()));
                return TError();
            }));

        TRequester requester(moduleMock, origDescr.ConnDescr());
        TChunksInput in{TChunkList{REQUEST_BODY}};
        TChunksOutput out;

        TryRethrowError(requester.TcpRequest(in, out));
        UNIT_ASSERT_EQUAL(out.Chunks(), RESPONSE_BODY);
    }

    Y_UNIT_TEST(AsyncRequesterTcpRequest) {
        TLog errorLog;
        NSrvKernel::NTesting::TModuleMock moduleMock;
        NSrvKernel::NTesting::TProcessMock process;

        ON_CALL(moduleMock, DoExtraAccessLog())
            .WillByDefault(Return(false));

        EXPECT_CALL(moduleMock, DoRun(_))
            .Times(1)
            .WillOnce(Invoke([&](const TConnDescr& descr) {
                UNIT_ASSERT_EQUAL(descr.ErrorLog, &errorLog);
                UNIT_ASSERT_EQUAL(descr.Request, nullptr);

                TChunkList chunk;
                TryRethrowError(descr.Input->Recv(chunk, TInstant::Max()));
                UNIT_ASSERT_EQUAL(chunk, REQUEST_BODY);

                TryRethrowError(descr.Output->Send(TChunkList(RESPONSE_BODY), TInstant::Max()));
                return TError();
            }));

        TAsyncRequester requester(moduleMock, &errorLog, process);
        TChunksInput in{TChunkList{REQUEST_BODY}};
        TChunksOutput out;

        TryRethrowError(requester.Requester().TcpRequest(in, out));
        UNIT_ASSERT_EQUAL(out.Chunks(), RESPONSE_BODY);
    }
}
